Skip to content

Commit 8c724db

Browse files
david-leifkerdavid-leifker
authored
feat(api): authorization extended for soft-delete and suspend (#12158)
1 parent 01a2c0c commit 8c724db

File tree

78 files changed

+980
-431
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

78 files changed

+980
-431
lines changed

‎datahub-frontend/app/auth/AuthModule.java

+2
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import io.datahubproject.metadata.context.EntityRegistryContext;
2828
import io.datahubproject.metadata.context.OperationContext;
2929
import io.datahubproject.metadata.context.OperationContextConfig;
30+
import io.datahubproject.metadata.context.RetrieverContext;
3031
import io.datahubproject.metadata.context.SearchContext;
3132
import io.datahubproject.metadata.context.ValidationContext;
3233
import java.nio.charset.StandardCharsets;
@@ -195,6 +196,7 @@ protected OperationContext provideOperationContext(
195196
.searchContext(SearchContext.EMPTY)
196197
.entityRegistryContext(EntityRegistryContext.builder().build(EmptyEntityRegistry.EMPTY))
197198
.validationContext(ValidationContext.builder().alternateValidation(false).build())
199+
.retrieverContext(RetrieverContext.EMPTY)
198200
.build(systemAuthentication);
199201
}
200202

‎datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/SystemUpdateConfig.java

+2
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import com.linkedin.gms.factory.kafka.common.TopicConventionFactory;
1414
import com.linkedin.gms.factory.kafka.schemaregistry.InternalSchemaRegistryFactory;
1515
import com.linkedin.gms.factory.search.BaseElasticSearchComponentsFactory;
16+
import com.linkedin.metadata.aspect.CachingAspectRetriever;
1617
import com.linkedin.metadata.config.kafka.KafkaConfiguration;
1718
import com.linkedin.metadata.dao.producer.KafkaEventProducer;
1819
import com.linkedin.metadata.dao.producer.KafkaHealthChecker;
@@ -186,6 +187,7 @@ protected OperationContext javaSystemOperationContext(
186187
components.getIndexConvention(),
187188
RetrieverContext.builder()
188189
.aspectRetriever(entityServiceAspectRetriever)
190+
.cachingAspectRetriever(CachingAspectRetriever.EMPTY)
189191
.graphRetriever(systemGraphRetriever)
190192
.searchRetriever(searchServiceSearchRetriever)
191193
.build(),

‎datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restorebackup/RestoreStorageStep.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,7 @@ private void readerExecutable(ReaderWrapper reader, UpgradeContext context) {
180180
try {
181181
aspectRecord =
182182
EntityUtils.toSystemAspect(
183-
context.opContext().getRetrieverContext().get(), aspect.toEntityAspect())
183+
context.opContext().getRetrieverContext(), aspect.toEntityAspect())
184184
.get()
185185
.getRecordTemplate();
186186
} catch (Exception e) {

‎datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/AbstractMCLStep.java

+1-2
Original file line numberDiff line numberDiff line change
@@ -113,8 +113,7 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
113113
List<Pair<Future<?>, SystemAspect>> futures;
114114
futures =
115115
EntityUtils.toSystemAspectFromEbeanAspects(
116-
opContext.getRetrieverContext().get(),
117-
batch.collect(Collectors.toList()))
116+
opContext.getRetrieverContext(), batch.collect(Collectors.toList()))
118117
.stream()
119118
.map(
120119
systemAspect -> {

‎datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/bootstrapmcps/BootstrapMCPUtil.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,8 @@ static AspectsBatch generateAspectBatch(
100100
.collect(Collectors.toList());
101101

102102
return AspectsBatchImpl.builder()
103-
.mcps(mcps, auditStamp, opContext.getRetrieverContext().get())
104-
.retrieverContext(opContext.getRetrieverContext().get())
103+
.mcps(mcps, auditStamp, opContext.getRetrieverContext())
104+
.retrieverContext(opContext.getRetrieverContext())
105105
.build();
106106
}
107107

‎datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/system/schemafield/GenerateSchemaFieldsFromSchemaMetadataStep.java

+3-7
Original file line numberDiff line numberDiff line change
@@ -168,13 +168,13 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
168168

169169
AspectsBatch aspectsBatch =
170170
AspectsBatchImpl.builder()
171-
.retrieverContext(opContext.getRetrieverContext().get())
171+
.retrieverContext(opContext.getRetrieverContext())
172172
.items(
173173
batch
174174
.flatMap(
175175
ebeanAspectV2 ->
176176
EntityUtils.toSystemAspectFromEbeanAspects(
177-
opContext.getRetrieverContext().get(),
177+
opContext.getRetrieverContext(),
178178
Set.of(ebeanAspectV2))
179179
.stream())
180180
.map(
@@ -189,11 +189,7 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
189189
.auditStamp(systemAspect.getAuditStamp())
190190
.systemMetadata(
191191
withAppSource(systemAspect.getSystemMetadata()))
192-
.build(
193-
opContext
194-
.getRetrieverContext()
195-
.get()
196-
.getAspectRetriever()))
192+
.build(opContext.getAspectRetriever()))
197193
.collect(Collectors.toList()))
198194
.build();
199195

‎datahub-upgrade/src/test/java/com/linkedin/datahub/upgrade/schemafield/GenerateSchemaFieldsFromSchemaMetadataStepTest.java

+1-2
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222
import com.linkedin.upgrade.DataHubUpgradeState;
2323
import io.datahubproject.metadata.context.OperationContext;
2424
import io.datahubproject.metadata.context.RetrieverContext;
25-
import java.util.Optional;
2625
import java.util.stream.Stream;
2726
import org.junit.jupiter.api.BeforeEach;
2827
import org.junit.jupiter.api.Test;
@@ -48,7 +47,7 @@ public void setup() {
4847
step =
4948
new GenerateSchemaFieldsFromSchemaMetadataStep(
5049
mockOpContext, mockEntityService, mockAspectDao, 10, 100, 1000);
51-
when(mockOpContext.getRetrieverContext()).thenReturn(Optional.of(mockRetrieverContext));
50+
when(mockOpContext.getRetrieverContext()).thenReturn(mockRetrieverContext);
5251
}
5352

5453
/** Test to verify the correct step ID is returned. */

‎entity-registry/src/main/java/com/linkedin/metadata/aspect/CachingAspectRetriever.java

+35-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,38 @@
11
package com.linkedin.metadata.aspect;
22

3+
import com.linkedin.common.urn.Urn;
4+
import com.linkedin.entity.Aspect;
5+
import com.linkedin.metadata.models.registry.EmptyEntityRegistry;
6+
import com.linkedin.metadata.models.registry.EntityRegistry;
7+
import java.util.Collections;
8+
import java.util.Map;
9+
import java.util.Set;
10+
import javax.annotation.Nonnull;
11+
312
/** Responses can be cached based on application.yaml caching configuration for the EntityClient */
4-
public interface CachingAspectRetriever extends AspectRetriever {}
13+
public interface CachingAspectRetriever extends AspectRetriever {
14+
15+
CachingAspectRetriever EMPTY = new EmptyAspectRetriever();
16+
17+
class EmptyAspectRetriever implements CachingAspectRetriever {
18+
@Nonnull
19+
@Override
20+
public Map<Urn, Map<String, Aspect>> getLatestAspectObjects(
21+
Set<Urn> urns, Set<String> aspectNames) {
22+
return Collections.emptyMap();
23+
}
24+
25+
@Nonnull
26+
@Override
27+
public Map<Urn, Map<String, SystemAspect>> getLatestSystemAspects(
28+
Map<Urn, Set<String>> urnAspectNames) {
29+
return Collections.emptyMap();
30+
}
31+
32+
@Nonnull
33+
@Override
34+
public EntityRegistry getEntityRegistry() {
35+
return EmptyEntityRegistry.EMPTY;
36+
}
37+
}
38+
}

‎entity-registry/src/main/java/com/linkedin/metadata/aspect/GraphRetriever.java

+23
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import com.linkedin.metadata.query.filter.Filter;
55
import com.linkedin.metadata.query.filter.RelationshipFilter;
66
import com.linkedin.metadata.query.filter.SortCriterion;
7+
import java.util.Collections;
78
import java.util.List;
89
import java.util.function.Function;
910
import javax.annotation.Nonnull;
@@ -97,4 +98,26 @@ default void consumeRelatedEntities(
9798
}
9899
}
99100
}
101+
102+
GraphRetriever EMPTY = new EmptyGraphRetriever();
103+
104+
class EmptyGraphRetriever implements GraphRetriever {
105+
106+
@Nonnull
107+
@Override
108+
public RelatedEntitiesScrollResult scrollRelatedEntities(
109+
@Nullable List<String> sourceTypes,
110+
@Nonnull Filter sourceEntityFilter,
111+
@Nullable List<String> destinationTypes,
112+
@Nonnull Filter destinationEntityFilter,
113+
@Nonnull List<String> relationshipTypes,
114+
@Nonnull RelationshipFilter relationshipFilter,
115+
@Nonnull List<SortCriterion> sortCriterion,
116+
@Nullable String scrollId,
117+
int count,
118+
@Nullable Long startTimeMillis,
119+
@Nullable Long endTimeMillis) {
120+
return new RelatedEntitiesScrollResult(0, 0, null, Collections.emptyList());
121+
}
122+
}
100123
}

‎entity-registry/src/main/java/com/linkedin/metadata/entity/SearchRetriever.java

+19
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.linkedin.metadata.query.filter.Filter;
44
import com.linkedin.metadata.search.ScrollResult;
5+
import com.linkedin.metadata.search.SearchEntityArray;
56
import java.util.List;
67
import javax.annotation.Nonnull;
78
import javax.annotation.Nullable;
@@ -21,4 +22,22 @@ ScrollResult scroll(
2122
@Nullable Filter filters,
2223
@Nullable String scrollId,
2324
int count);
25+
26+
SearchRetriever EMPTY = new EmptySearchRetriever();
27+
28+
class EmptySearchRetriever implements SearchRetriever {
29+
30+
@Override
31+
public ScrollResult scroll(
32+
@Nonnull List<String> entities,
33+
@Nullable Filter filters,
34+
@Nullable String scrollId,
35+
int count) {
36+
ScrollResult empty = new ScrollResult();
37+
empty.setEntities(new SearchEntityArray());
38+
empty.setNumEntities(0);
39+
empty.setPageSize(0);
40+
return empty;
41+
}
42+
}
2443
}

‎entity-registry/src/testFixtures/java/com/linkedin/test/metadata/aspect/MockAspectRetriever.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import com.linkedin.data.DataMap;
66
import com.linkedin.data.template.RecordTemplate;
77
import com.linkedin.entity.Aspect;
8-
import com.linkedin.metadata.aspect.AspectRetriever;
8+
import com.linkedin.metadata.aspect.CachingAspectRetriever;
99
import com.linkedin.metadata.aspect.SystemAspect;
1010
import com.linkedin.metadata.models.registry.EntityRegistry;
1111
import com.linkedin.mxe.SystemMetadata;
@@ -22,7 +22,7 @@
2222
import javax.annotation.Nonnull;
2323
import org.mockito.Mockito;
2424

25-
public class MockAspectRetriever implements AspectRetriever {
25+
public class MockAspectRetriever implements CachingAspectRetriever {
2626
private final Map<Urn, Map<String, Aspect>> data;
2727
private final Map<Urn, Map<String, SystemAspect>> systemData = new HashMap<>();
2828

‎li-utils/src/main/java/com/linkedin/metadata/Constants.java

+2
Original file line numberDiff line numberDiff line change
@@ -409,6 +409,8 @@ public class Constants {
409409
/** User Status */
410410
public static final String CORP_USER_STATUS_ACTIVE = "ACTIVE";
411411

412+
public static final String CORP_USER_STATUS_SUSPENDED = "SUSPENDED";
413+
412414
/** Task Runs */
413415
public static final String DATA_PROCESS_INSTANCE_ENTITY_NAME = "dataProcessInstance";
414416

‎metadata-io/metadata-io-api/src/test/java/com/linkedin/metadata/entity/ebean/batch/AspectsBatchImplTest.java

+4-4
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
import com.linkedin.data.schema.annotation.PathSpecBasedSchemaAnnotationVisitor;
1717
import com.linkedin.dataset.DatasetProperties;
1818
import com.linkedin.events.metadata.ChangeType;
19-
import com.linkedin.metadata.aspect.AspectRetriever;
19+
import com.linkedin.metadata.aspect.CachingAspectRetriever;
2020
import com.linkedin.metadata.aspect.GraphRetriever;
2121
import com.linkedin.metadata.aspect.batch.MCPItem;
2222
import com.linkedin.metadata.aspect.patch.GenericJsonPatch;
@@ -56,7 +56,7 @@
5656

5757
public class AspectsBatchImplTest {
5858
private EntityRegistry testRegistry;
59-
private AspectRetriever mockAspectRetriever;
59+
private CachingAspectRetriever mockAspectRetriever;
6060
private RetrieverContext retrieverContext;
6161

6262
@BeforeTest
@@ -75,12 +75,12 @@ public void beforeTest() throws EntityRegistryException {
7575

7676
@BeforeMethod
7777
public void setup() {
78-
this.mockAspectRetriever = mock(AspectRetriever.class);
78+
this.mockAspectRetriever = mock(CachingAspectRetriever.class);
7979
when(this.mockAspectRetriever.getEntityRegistry()).thenReturn(testRegistry);
8080
this.retrieverContext =
8181
RetrieverContext.builder()
8282
.searchRetriever(mock(SearchRetriever.class))
83-
.aspectRetriever(mockAspectRetriever)
83+
.cachingAspectRetriever(mockAspectRetriever)
8484
.graphRetriever(mock(GraphRetriever.class))
8585
.build();
8686
}

‎metadata-io/src/main/java/com/linkedin/metadata/aspect/utils/DefaultAspectsUtil.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ public static List<MCPItem> getAdditionalChanges(
137137
getProposalFromAspectForDefault(
138138
entry.getKey(), entry.getValue(), entityKeyAspect, templateItem),
139139
templateItem.getAuditStamp(),
140-
opContext.getAspectRetrieverOpt().get()))
140+
opContext.getAspectRetriever()))
141141
.filter(Objects::nonNull);
142142
})
143143
.collect(Collectors.toList());

‎metadata-io/src/main/java/com/linkedin/metadata/client/EntityClientAspectRetriever.java

+4-3
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ public EntityRegistry getEntityRegistry() {
3535
@Override
3636
public Aspect getLatestAspectObject(@Nonnull Urn urn, @Nonnull String aspectName) {
3737
try {
38-
return entityClient.getLatestAspectObject(systemOperationContext, urn, aspectName);
38+
return entityClient.getLatestAspectObject(systemOperationContext, urn, aspectName, false);
3939
} catch (RemoteInvocationException | URISyntaxException e) {
4040
throw new RuntimeException(e);
4141
}
@@ -49,7 +49,7 @@ public Map<Urn, Map<String, Aspect>> getLatestAspectObjects(
4949
return Map.of();
5050
} else {
5151
try {
52-
return entityClient.getLatestAspects(systemOperationContext, urns, aspectNames);
52+
return entityClient.getLatestAspects(systemOperationContext, urns, aspectNames, false);
5353
} catch (RemoteInvocationException | URISyntaxException e) {
5454
throw new RuntimeException(e);
5555
}
@@ -70,7 +70,8 @@ public Map<Urn, Map<String, SystemAspect>> getLatestSystemAspects(
7070
urnAspectNames.keySet(),
7171
urnAspectNames.values().stream()
7272
.flatMap(Collection::stream)
73-
.collect(Collectors.toSet()));
73+
.collect(Collectors.toSet()),
74+
false);
7475
} catch (RemoteInvocationException | URISyntaxException e) {
7576
throw new RuntimeException(e);
7677
}

‎metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java

+16-5
Original file line numberDiff line numberDiff line change
@@ -106,11 +106,17 @@ public EntityResponse getV2(
106106
@Nonnull OperationContext opContext,
107107
@Nonnull String entityName,
108108
@Nonnull final Urn urn,
109-
@Nullable final Set<String> aspectNames)
109+
@Nullable final Set<String> aspectNames,
110+
@Nullable Boolean alwaysIncludeKeyAspect)
110111
throws RemoteInvocationException, URISyntaxException {
111112
final Set<String> projectedAspects =
112113
aspectNames == null ? opContext.getEntityAspectNames(entityName) : aspectNames;
113-
return entityService.getEntityV2(opContext, entityName, urn, projectedAspects);
114+
return entityService.getEntityV2(
115+
opContext,
116+
entityName,
117+
urn,
118+
projectedAspects,
119+
alwaysIncludeKeyAspect == null || alwaysIncludeKeyAspect);
114120
}
115121

116122
@Override
@@ -126,7 +132,8 @@ public Map<Urn, EntityResponse> batchGetV2(
126132
@Nonnull OperationContext opContext,
127133
@Nonnull String entityName,
128134
@Nonnull Set<Urn> urns,
129-
@Nullable Set<String> aspectNames)
135+
@Nullable Set<String> aspectNames,
136+
@Nullable Boolean alwaysIncludeKeyAspect)
130137
throws RemoteInvocationException, URISyntaxException {
131138
final Set<String> projectedAspects =
132139
aspectNames == null ? opContext.getEntityAspectNames(entityName) : aspectNames;
@@ -139,7 +146,11 @@ public Map<Urn, EntityResponse> batchGetV2(
139146
try {
140147
responseMap.putAll(
141148
entityService.getEntitiesV2(
142-
opContext, entityName, new HashSet<>(batch), projectedAspects));
149+
opContext,
150+
entityName,
151+
new HashSet<>(batch),
152+
projectedAspects,
153+
alwaysIncludeKeyAspect == null || alwaysIncludeKeyAspect));
143154
} catch (URISyntaxException e) {
144155
throw new RuntimeException(e);
145156
}
@@ -772,7 +783,7 @@ public List<String> batchIngestProposals(
772783
.mcps(
773784
batch,
774785
auditStamp,
775-
opContext.getRetrieverContext().get(),
786+
opContext.getRetrieverContext(),
776787
opContext.getValidationContext().isAlternateValidation())
777788
.build();
778789

‎metadata-io/src/main/java/com/linkedin/metadata/client/SystemJavaEntityClient.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,6 @@ public Map<Urn, EntityResponse> batchGetV2NoCache(
8989
@Nonnull Set<Urn> urns,
9090
@Nullable Set<String> aspectNames)
9191
throws RemoteInvocationException, URISyntaxException {
92-
return super.batchGetV2(opContext, entityName, urns, aspectNames);
92+
return super.batchGetV2(opContext, entityName, urns, aspectNames, false);
9393
}
9494
}

0 commit comments

Comments
 (0)