feat(graphql): extend entity client to support aspect methods directl…

…y via java (datahub-project#3489)

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

@@ -72,6 +72,8 @@ public static RestliEntityClient getEntitiesClient() {
         return _entities;
     }
 
+    // Deprecated- please use EntityClient from now on for all aspect related calls
+    @Deprecated
     public static AspectClient getAspectsClient() {
         if (_aspects == null) {
             synchronized (GmsClientFactory.class) {

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

@@ -235,7 +235,7 @@ public GmsGraphQLEngine(
             GmsClientFactory.getRelationshipsClient()
         );
         this.glossaryTermType = new GlossaryTermType(entityClient);
-        this.aspectType = new AspectType(GmsClientFactory.getAspectsClient());
+        this.aspectType = new AspectType(entityClient);
         this.usageType = new UsageType(GmsClientFactory.getUsageClient());
 
         // Init Lists

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/aspect/AspectType.java

@@ -3,7 +3,7 @@
 import com.linkedin.datahub.graphql.VersionedAspectKey;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.Aspect;
-import com.linkedin.entity.client.AspectClient;
+import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.aspect.VersionedAspect;
 import com.linkedin.r2.RemoteInvocationException;
 import com.linkedin.restli.client.RestLiResponseException;
@@ -14,10 +14,10 @@
 
 
 public class AspectType {
-  private final AspectClient _aspectClient;
+  private final EntityClient _entityClient;
 
-  public AspectType(final AspectClient aspectClient) {
-    _aspectClient = aspectClient;
+  public AspectType(final EntityClient entityClient) {
+    _entityClient = entityClient;
   }
 
   /**
@@ -30,7 +30,7 @@ public List<DataFetcherResult<Aspect>> batchLoad(@Nonnull List<VersionedAspectKe
     try {
       return keys.stream().map(key -> {
         try {
-          VersionedAspect entity = _aspectClient.getAspect(key.getUrn(), key.getAspectName(), key.getVersion(), context.getActor());
+          VersionedAspect entity = _entityClient.getAspect(key.getUrn(), key.getAspectName(), key.getVersion(), context.getActor());
           return DataFetcherResult.<Aspect>newResult().data(AspectMapper.map(entity)).build();
         } catch (RemoteInvocationException e) {
           if (e instanceof RestLiResponseException) {

metadata-service/graphql-api/src/main/java/com/datahub/metadata/graphql/JavaEntityClientFactory.java

@@ -4,6 +4,7 @@
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.search.EntitySearchService;
 import com.linkedin.metadata.search.SearchService;
+import com.linkedin.metadata.timeseries.TimeseriesAspectService;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.context.annotation.Bean;
@@ -24,8 +25,12 @@ public class JavaEntityClientFactory {
   @Qualifier("entitySearchService")
   private EntitySearchService _entitySearchService;
 
+  @Autowired
+  @Qualifier("timeseriesAspectService")
+  private TimeseriesAspectService _timeseriesAspectService;
+
   @Bean("javaEntityClient")
   public JavaEntityClient getJavaEntityClient() {
-    return new JavaEntityClient(_entityService, _entitySearchService, _searchService);
+    return new JavaEntityClient(_entityService, _entitySearchService, _searchService, _timeseriesAspectService);
   }
 }

metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java

@@ -1,19 +1,24 @@
 package com.linkedin.entity.client;
 
 import com.linkedin.common.urn.Urn;
+import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.data.template.StringArray;
 import com.linkedin.entity.Entity;
+import com.linkedin.metadata.aspect.EnvelopedAspect;
+import com.linkedin.metadata.aspect.VersionedAspect;
 import com.linkedin.metadata.browse.BrowseResult;
 import com.linkedin.metadata.query.AutoCompleteResult;
 import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.query.ListResult;
 import com.linkedin.metadata.query.filter.SortCriterion;
 import com.linkedin.metadata.search.SearchResult;
 import com.linkedin.metadata.query.ListUrnsResult;
+import com.linkedin.mxe.MetadataChangeProposal;
 import com.linkedin.mxe.SystemMetadata;
 import com.linkedin.r2.RemoteInvocationException;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -218,4 +223,39 @@ public SearchResult filter(
       int count,
       @Nonnull String actor)
       throws RemoteInvocationException;
+
+  public VersionedAspect getAspect(
+      @Nonnull String urn,
+      @Nonnull String aspect,
+      @Nonnull Long version,
+      @Nonnull String actor)
+      throws RemoteInvocationException;
+
+  public VersionedAspect getAspectOrNull(
+      @Nonnull String urn,
+      @Nonnull String aspect,
+      @Nonnull Long version,
+      @Nonnull String actor) throws RemoteInvocationException;
+
+  public List<EnvelopedAspect> getTimeseriesAspectValues(
+      @Nonnull String urn,
+      @Nonnull String entity,
+      @Nonnull String aspect,
+      @Nullable Long startTimeMillis,
+      @Nullable Long endTimeMillis,
+      @Nullable Integer limit,
+      @Nullable String actor
+  ) throws RemoteInvocationException;
+
+  public String ingestProposal(
+      @Nonnull final MetadataChangeProposal metadataChangeProposal,
+      @Nonnull final String actor
+  ) throws RemoteInvocationException;
+
+  public <T extends RecordTemplate> Optional<T> getVersionedAspect(
+      @Nonnull String urn,
+      @Nonnull String aspect,
+      @Nonnull Long version,
+      @Nonnull String actor,
+      @Nonnull Class<T> aspectClass) throws RemoteInvocationException;
 }

metadata-service/restli-client/src/main/java/com/linkedin/entity/client/JavaEntityClient.java

@@ -1,48 +1,67 @@
 package com.linkedin.entity.client;
 
 import com.google.common.collect.ImmutableList;
+
 import com.google.common.collect.ImmutableSet;
+import com.linkedin.aspect.GetTimeseriesAspectValuesResponse;
 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.StringArray;
 import com.linkedin.entity.Entity;
+import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.aspect.EnvelopedAspect;
+import com.linkedin.metadata.aspect.EnvelopedAspectArray;
+import com.linkedin.metadata.aspect.VersionedAspect;
 import com.linkedin.metadata.browse.BrowseResult;
+import com.linkedin.metadata.dao.utils.RecordUtils;
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.query.AutoCompleteResult;
 import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.query.filter.SortCriterion;
 import com.linkedin.metadata.query.ListResult;
 import com.linkedin.metadata.query.ListUrnsResult;
+import com.linkedin.metadata.resources.entity.AspectUtils;
 import com.linkedin.metadata.resources.entity.EntityResource;
 import com.linkedin.metadata.search.EntitySearchService;
 import com.linkedin.metadata.search.SearchResult;
 import com.linkedin.metadata.search.SearchService;
+import com.linkedin.metadata.timeseries.TimeseriesAspectService;
+import com.linkedin.mxe.MetadataChangeProposal;
 import com.linkedin.mxe.SystemMetadata;
 import com.linkedin.r2.RemoteInvocationException;
 import java.time.Clock;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import lombok.SneakyThrows;
+import lombok.extern.slf4j.Slf4j;
 
 import static com.linkedin.metadata.search.utils.QueryUtils.*;
 
 
+@Slf4j
 public class JavaEntityClient implements EntityClient {
 
+    private final Clock _clock = Clock.systemUTC();
+
     private EntityService _entityService;
     private EntitySearchService _entitySearchService;
     private SearchService _searchService;
+    private TimeseriesAspectService _timeseriesAspectService;
 
     public JavaEntityClient(@Nonnull final EntityService entityService, @Nonnull final EntitySearchService entitySearchService, @Nonnull final
-        SearchService searchService) {
+        SearchService searchService, @Nonnull final TimeseriesAspectService timeseriesAspectService) {
       _entityService = entityService;
       _entitySearchService = entitySearchService;
       _searchService = searchService;
+      _timeseriesAspectService = timeseriesAspectService;
     }
 
     @Nonnull
@@ -285,4 +304,69 @@ public SearchResult filter(@Nonnull String entity, @Nonnull Filter filter, @Null
         int start, int count, @Nonnull String actor) throws RemoteInvocationException {
         return _entitySearchService.filter(entity, filter, sortCriterion, start, count);
     }
+
+    @SneakyThrows
+    @Override
+    public VersionedAspect getAspect(@Nonnull String urn, @Nonnull String aspect, @Nonnull Long version,
+        @Nonnull String actor) throws RemoteInvocationException {
+        return _entityService.getVersionedAspect(Urn.createFromString(urn), aspect, version);
+    }
+
+    @SneakyThrows
+    @Override
+    public VersionedAspect getAspectOrNull(@Nonnull String urn, @Nonnull String aspect, @Nonnull Long version,
+        @Nonnull String actor) throws RemoteInvocationException {
+        return _entityService.getVersionedAspect(Urn.createFromString(urn), aspect, version);
+    }
+
+    @SneakyThrows
+    @Override
+    public List<EnvelopedAspect> getTimeseriesAspectValues(@Nonnull String urn, @Nonnull String entity,
+        @Nonnull String aspect, @Nullable Long startTimeMillis, @Nullable Long endTimeMillis, @Nullable Integer limit,
+        @Nullable String actor) throws RemoteInvocationException {
+        GetTimeseriesAspectValuesResponse response = new GetTimeseriesAspectValuesResponse();
+        response.setEntityName(entity);
+        response.setAspectName(aspect);
+        if (startTimeMillis != null) {
+            response.setStartTimeMillis(startTimeMillis);
+        }
+        if (endTimeMillis != null) {
+            response.setEndTimeMillis(endTimeMillis);
+        }
+        response.setLimit(limit);
+        response.setValues(new EnvelopedAspectArray(
+            _timeseriesAspectService.getAspectValues(Urn.createFromString(urn), entity, aspect, startTimeMillis, endTimeMillis,
+                limit)));
+        return response.getValues();
+    }
+
+    // TODO: Factor out ingest logic into a util that can be accessed by the java client and the resource
+    @SneakyThrows
+    @Override
+    public String ingestProposal(@Nonnull MetadataChangeProposal metadataChangeProposal,
+        @Nonnull String actor) throws RemoteInvocationException {
+        final AuditStamp auditStamp =
+            new AuditStamp().setTime(_clock.millis()).setActor(Urn.createFromString(Constants.UNKNOWN_ACTOR));
+        final List<MetadataChangeProposal> additionalChanges =
+            AspectUtils.getAdditionalChanges(metadataChangeProposal, _entityService);
+
+        Urn urn = _entityService.ingestProposal(metadataChangeProposal, auditStamp);
+        additionalChanges.forEach(proposal -> _entityService.ingestProposal(proposal, auditStamp));
+        return urn.toString();
+    }
+
+    @SneakyThrows
+    @Override
+    public <T extends RecordTemplate> Optional<T> getVersionedAspect(@Nonnull String urn, @Nonnull String aspect,
+        @Nonnull Long version, @Nonnull String actor, @Nonnull Class<T> aspectClass) throws RemoteInvocationException {
+        VersionedAspect entity = _entityService.getVersionedAspect(Urn.createFromString(urn), aspect, version);
+        if (entity.hasAspect()) {
+            DataMap rawAspect = ((DataMap) entity.data().get("aspect"));
+            if (rawAspect.containsKey(aspectClass.getCanonicalName())) {
+                DataMap aspectDataMap = rawAspect.getDataMap(aspectClass.getCanonicalName());
+                return Optional.of(RecordUtils.toRecordTemplate(aspectClass, aspectDataMap));
+            }
+        }
+        return Optional.empty();
+    }
 }