feat(auth): add data platform instance field resolver provider (#8828)

Co-authored-by: Sergio Gómez Villamor <sgomezvillamor@gmail.com>
Co-authored-by: Adrián Pertíñez <khurzak92@gmail.com>

metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResolvedResourceSpec.java

@@ -3,6 +3,7 @@ package com.datahub.authorization;
 import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
+import javax.annotation.Nullable;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import lombok.ToString;
@@ -35,4 +36,20 @@ public class ResolvedResourceSpec {
     }
     return fieldResolvers.get(ResourceFieldType.OWNER).getFieldValuesFuture().join().getValues();
   }
+
+  /**
+   * Fetch the platform instance for a Resolved Resource Spec
+   * @return a Platform Instance or null if one does not exist.
+   */
+  @Nullable
+  public String getDataPlatformInstance() {
+    if (!fieldResolvers.containsKey(ResourceFieldType.DATA_PLATFORM_INSTANCE)) {
+      return null;
+    }
+    Set<String> dataPlatformInstance = fieldResolvers.get(ResourceFieldType.DATA_PLATFORM_INSTANCE).getFieldValuesFuture().join().getValues();
+    if (dataPlatformInstance.size() > 0) {
+      return dataPlatformInstance.stream().findFirst().get();
+    }
+    return null;
+  }
 }

metadata-auth/auth-api/src/main/java/com/datahub/authorization/ResourceFieldType.java

@@ -19,5 +19,9 @@ public enum ResourceFieldType {
   /**
    * Domains of resource
    */
-  DOMAIN
+  DOMAIN,
+  /**
+   * Data platform instance of resource
+   */
+  DATA_PLATFORM_INSTANCE
 }

metadata-service/auth-impl/src/main/java/com/datahub/authorization/DefaultResourceSpecResolver.java

@@ -1,13 +1,15 @@
 package com.datahub.authorization;
 
-import com.datahub.authorization.fieldresolverprovider.EntityTypeFieldResolverProvider;
-import com.datahub.authorization.fieldresolverprovider.OwnerFieldResolverProvider;
 import com.datahub.authentication.Authentication;
+import com.datahub.authorization.fieldresolverprovider.DataPlatformInstanceFieldResolverProvider;
 import com.datahub.authorization.fieldresolverprovider.DomainFieldResolverProvider;
+import com.datahub.authorization.fieldresolverprovider.EntityTypeFieldResolverProvider;
 import com.datahub.authorization.fieldresolverprovider.EntityUrnFieldResolverProvider;
+import com.datahub.authorization.fieldresolverprovider.OwnerFieldResolverProvider;
 import com.datahub.authorization.fieldresolverprovider.ResourceFieldResolverProvider;
 import com.google.common.collect.ImmutableList;
 import com.linkedin.entity.client.EntityClient;
+
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -20,7 +22,8 @@ public class DefaultResourceSpecResolver implements ResourceSpecResolver {
     _resourceFieldResolverProviders =
         ImmutableList.of(new EntityTypeFieldResolverProvider(), new EntityUrnFieldResolverProvider(),
             new DomainFieldResolverProvider(entityClient, systemAuthentication),
-            new OwnerFieldResolverProvider(entityClient, systemAuthentication));
+            new OwnerFieldResolverProvider(entityClient, systemAuthentication),
+            new DataPlatformInstanceFieldResolverProvider(entityClient, systemAuthentication));
   }
 
   @Override

metadata-service/auth-impl/src/main/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProvider.java

@@ -0,0 +1,70 @@
+package com.datahub.authorization.fieldresolverprovider;
+
+import com.datahub.authentication.Authentication;
+import com.datahub.authorization.FieldResolver;
+import com.datahub.authorization.ResourceFieldType;
+import com.datahub.authorization.ResourceSpec;
+import com.linkedin.common.DataPlatformInstance;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.EnvelopedAspect;
+import com.linkedin.entity.client.EntityClient;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.Collections;
+import java.util.Objects;
+
+import static com.linkedin.metadata.Constants.*;
+
+/**
+ * Provides field resolver for domain given resourceSpec
+ */
+@Slf4j
+@RequiredArgsConstructor
+public class DataPlatformInstanceFieldResolverProvider implements ResourceFieldResolverProvider {
+
+  private final EntityClient _entityClient;
+  private final Authentication _systemAuthentication;
+
+  @Override
+  public ResourceFieldType getFieldType() {
+    return ResourceFieldType.DATA_PLATFORM_INSTANCE;
+  }
+
+  @Override
+  public FieldResolver getFieldResolver(ResourceSpec resourceSpec) {
+    return FieldResolver.getResolverFromFunction(resourceSpec, this::getDataPlatformInstance);
+  }
+
+  private FieldResolver.FieldValue getDataPlatformInstance(ResourceSpec resourceSpec) {
+    Urn entityUrn = UrnUtils.getUrn(resourceSpec.getResource());
+    // In the case that the entity is a platform instance, the associated platform instance entity is the instance itself
+    if (entityUrn.getEntityType().equals(DATA_PLATFORM_INSTANCE_ENTITY_NAME)) {
+      return FieldResolver.FieldValue.builder()
+          .values(Collections.singleton(entityUrn.toString()))
+          .build();
+    }
+
+    EnvelopedAspect dataPlatformInstanceAspect;
+    try {
+      EntityResponse response = _entityClient.getV2(entityUrn.getEntityType(), entityUrn,
+          Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME), _systemAuthentication);
+      if (response == null || !response.getAspects().containsKey(DATA_PLATFORM_INSTANCE_ASPECT_NAME)) {
+        return FieldResolver.emptyFieldValue();
+      }
+      dataPlatformInstanceAspect = response.getAspects().get(DATA_PLATFORM_INSTANCE_ASPECT_NAME);
+    } catch (Exception e) {
+      log.error("Error while retrieving platform instance aspect for urn {}", entityUrn, e);
+      return FieldResolver.emptyFieldValue();
+    }
+    DataPlatformInstance dataPlatformInstance = new DataPlatformInstance(dataPlatformInstanceAspect.getValue().data());
+    if (dataPlatformInstance.getInstance() == null) {
+      return FieldResolver.emptyFieldValue();
+    }
+    return FieldResolver.FieldValue.builder()
+        .values(Collections.singleton(Objects.requireNonNull(dataPlatformInstance.getInstance()).toString()))
+        .build();
+  }
+}

metadata-service/auth-impl/src/test/java/com/datahub/authorization/fieldresolverprovider/DataPlatformInstanceFieldResolverProviderTest.java

@@ -0,0 +1,188 @@
+package com.datahub.authorization.fieldresolverprovider;
+
+import com.datahub.authentication.Authentication;
+import com.datahub.authorization.ResourceFieldType;
+import com.datahub.authorization.ResourceSpec;
+import com.linkedin.common.DataPlatformInstance;
+import com.linkedin.common.urn.Urn;
+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.r2.RemoteInvocationException;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import java.net.URISyntaxException;
+import java.util.Collections;
+import java.util.Set;
+
+import static com.linkedin.metadata.Constants.*;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.*;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+public class DataPlatformInstanceFieldResolverProviderTest {
+
+  private static final String DATA_PLATFORM_INSTANCE_URN =
+      "urn:li:dataPlatformInstance:(urn:li:dataPlatform:s3,test-platform-instance)";
+  private static final String RESOURCE_URN =
+      "urn:li:dataset:(urn:li:dataPlatform:s3,test-platform-instance.testDataset,PROD)";
+  private static final ResourceSpec RESOURCE_SPEC = new ResourceSpec(DATASET_ENTITY_NAME, RESOURCE_URN);
+
+  @Mock
+  private EntityClient entityClientMock;
+  @Mock
+  private Authentication systemAuthenticationMock;
+
+  private DataPlatformInstanceFieldResolverProvider dataPlatformInstanceFieldResolverProvider;
+
+  @BeforeMethod
+  public void setup() {
+    MockitoAnnotations.initMocks(this);
+    dataPlatformInstanceFieldResolverProvider =
+        new DataPlatformInstanceFieldResolverProvider(entityClientMock, systemAuthenticationMock);
+  }
+
+  @Test
+  public void shouldReturnDataPlatformInstanceType() {
+    assertEquals(ResourceFieldType.DATA_PLATFORM_INSTANCE, dataPlatformInstanceFieldResolverProvider.getFieldType());
+  }
+
+  @Test
+  public void shouldReturnFieldValueWithResourceSpecIfTypeIsDataPlatformInstance() {
+    var resourceSpec = new ResourceSpec(DATA_PLATFORM_INSTANCE_ENTITY_NAME, DATA_PLATFORM_INSTANCE_URN);
+
+    var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(resourceSpec);
+
+    assertEquals(Set.of(DATA_PLATFORM_INSTANCE_URN), result.getFieldValuesFuture().join().getValues());
+    verifyZeroInteractions(entityClientMock);
+  }
+
+  @Test
+  public void shouldReturnEmptyFieldValueWhenResponseIsNull() throws RemoteInvocationException, URISyntaxException {
+    when(entityClientMock.getV2(
+        eq(DATASET_ENTITY_NAME),
+        any(Urn.class),
+        eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)),
+        eq(systemAuthenticationMock)
+    )).thenReturn(null);
+
+    var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(RESOURCE_SPEC);
+
+    assertTrue(result.getFieldValuesFuture().join().getValues().isEmpty());
+    verify(entityClientMock, times(1)).getV2(
+        eq(DATASET_ENTITY_NAME),
+        any(Urn.class),
+        eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)),
+        eq(systemAuthenticationMock)
+    );
+  }
+
+  @Test
+  public void shouldReturnEmptyFieldValueWhenResourceHasNoDataPlatformInstance()
+      throws RemoteInvocationException, URISyntaxException {
+    var entityResponseMock = mock(EntityResponse.class);
+    when(entityResponseMock.getAspects()).thenReturn(new EnvelopedAspectMap());
+    when(entityClientMock.getV2(
+        eq(DATASET_ENTITY_NAME),
+        any(Urn.class),
+        eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)),
+        eq(systemAuthenticationMock)
+    )).thenReturn(entityResponseMock);
+
+    var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(RESOURCE_SPEC);
+
+    assertTrue(result.getFieldValuesFuture().join().getValues().isEmpty());
+    verify(entityClientMock, times(1)).getV2(
+        eq(DATASET_ENTITY_NAME),
+        any(Urn.class),
+        eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)),
+        eq(systemAuthenticationMock)
+    );
+  }
+
+  @Test
+  public void shouldReturnEmptyFieldValueWhenThereIsAnException() throws RemoteInvocationException, URISyntaxException {
+    when(entityClientMock.getV2(
+        eq(DATASET_ENTITY_NAME),
+        any(Urn.class),
+        eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)),
+        eq(systemAuthenticationMock)
+    )).thenThrow(new RemoteInvocationException());
+
+    var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(RESOURCE_SPEC);
+
+    assertTrue(result.getFieldValuesFuture().join().getValues().isEmpty());
+    verify(entityClientMock, times(1)).getV2(
+        eq(DATASET_ENTITY_NAME),
+        any(Urn.class),
+        eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)),
+        eq(systemAuthenticationMock)
+    );
+  }
+
+  @Test
+  public void shouldReturnEmptyFieldValueWhenDataPlatformInstanceHasNoInstance()
+      throws RemoteInvocationException, URISyntaxException {
+
+    var dataPlatform = new DataPlatformInstance()
+        .setPlatform(Urn.createFromString("urn:li:dataPlatform:s3"));
+    var entityResponseMock = mock(EntityResponse.class);
+    var envelopedAspectMap = new EnvelopedAspectMap();
+    envelopedAspectMap.put(DATA_PLATFORM_INSTANCE_ASPECT_NAME,
+        new EnvelopedAspect().setValue(new Aspect(dataPlatform.data())));
+    when(entityResponseMock.getAspects()).thenReturn(envelopedAspectMap);
+    when(entityClientMock.getV2(
+        eq(DATASET_ENTITY_NAME),
+        any(Urn.class),
+        eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)),
+        eq(systemAuthenticationMock)
+    )).thenReturn(entityResponseMock);
+
+    var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(RESOURCE_SPEC);
+
+    assertTrue(result.getFieldValuesFuture().join().getValues().isEmpty());
+    verify(entityClientMock, times(1)).getV2(
+        eq(DATASET_ENTITY_NAME),
+        any(Urn.class),
+        eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)),
+        eq(systemAuthenticationMock)
+    );
+  }
+
+  @Test
+  public void shouldReturnFieldValueWithDataPlatformInstanceOfTheResource()
+      throws RemoteInvocationException, URISyntaxException {
+
+    var dataPlatformInstance = new DataPlatformInstance()
+        .setPlatform(Urn.createFromString("urn:li:dataPlatform:s3"))
+        .setInstance(Urn.createFromString(DATA_PLATFORM_INSTANCE_URN));
+    var entityResponseMock = mock(EntityResponse.class);
+    var envelopedAspectMap = new EnvelopedAspectMap();
+    envelopedAspectMap.put(DATA_PLATFORM_INSTANCE_ASPECT_NAME,
+        new EnvelopedAspect().setValue(new Aspect(dataPlatformInstance.data())));
+    when(entityResponseMock.getAspects()).thenReturn(envelopedAspectMap);
+    when(entityClientMock.getV2(
+        eq(DATASET_ENTITY_NAME),
+        any(Urn.class),
+        eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)),
+        eq(systemAuthenticationMock)
+    )).thenReturn(entityResponseMock);
+
+    var result = dataPlatformInstanceFieldResolverProvider.getFieldResolver(RESOURCE_SPEC);
+
+    assertEquals(Set.of(DATA_PLATFORM_INSTANCE_URN), result.getFieldValuesFuture().join().getValues());
+    verify(entityClientMock, times(1)).getV2(
+        eq(DATASET_ENTITY_NAME),
+        any(Urn.class),
+        eq(Collections.singleton(DATA_PLATFORM_INSTANCE_ASPECT_NAME)),
+        eq(systemAuthenticationMock)
+    );
+  }
+}