fix(upgrade): Improving NoCodeUpgrade logic to account for Bootstrap logic (#3301)

This commit is contained in:
John Joyce 2021-09-28 16:30:49 -07:00 committed by GitHub
parent add778c04a
commit 33f4d2ede0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
16 changed files with 37 additions and 98 deletions

View File

@ -15,6 +15,7 @@ import com.linkedin.identity.CorpGroupInfo;
import com.linkedin.identity.CorpUserEditableInfo;
import com.linkedin.identity.CorpUserInfo;
import com.linkedin.identity.GroupMembership;
import com.linkedin.metadata.Constants;
import com.linkedin.metadata.aspect.CorpGroupAspect;
import com.linkedin.metadata.aspect.CorpGroupAspectArray;
import com.linkedin.metadata.aspect.CorpUserAspect;
@ -64,7 +65,7 @@ import static auth.AuthUtils.*;
@Slf4j
public class OidcCallbackLogic extends DefaultCallbackLogic<Result, PlayWebContext> {
private static final String SYSTEM_PRINCIPAL = "urn:li:corpuser:system";
private static final String SYSTEM_ACTOR = Constants.SYSTEM_ACTOR;
private final EntityClient _entityClient = GmsClientFactory.getEntitiesClient();
private final SsoManager _ssoManager;
@ -271,7 +272,7 @@ public class OidcCallbackLogic extends DefaultCallbackLogic<Result, PlayWebConte
// 1. Check if this user already exists.
try {
final Entity corpUser = _entityClient.get(corpUserSnapshot.getUrn(), SYSTEM_PRINCIPAL);
final Entity corpUser = _entityClient.get(corpUserSnapshot.getUrn(), SYSTEM_ACTOR);
log.debug(String.format("Fetched GMS user with urn %s",corpUserSnapshot.getUrn()));
@ -281,7 +282,7 @@ public class OidcCallbackLogic extends DefaultCallbackLogic<Result, PlayWebConte
// 2. The user does not exist. Provision them.
final Entity newEntity = new Entity();
newEntity.setValue(Snapshot.create(corpUserSnapshot));
_entityClient.update(newEntity, SYSTEM_PRINCIPAL);
_entityClient.update(newEntity, SYSTEM_ACTOR);
log.debug(String.format("Successfully provisioned user %s", corpUserSnapshot.getUrn()));
}
@ -302,7 +303,7 @@ public class OidcCallbackLogic extends DefaultCallbackLogic<Result, PlayWebConte
// 1. Check if this user already exists.
try {
final Set<Urn> urnsToFetch = corpGroups.stream().map(CorpGroupSnapshot::getUrn).collect(Collectors.toSet());
final Map<Urn, Entity> existingGroups = _entityClient.batchGet(urnsToFetch, SYSTEM_PRINCIPAL);
final Map<Urn, Entity> existingGroups = _entityClient.batchGet(urnsToFetch, SYSTEM_ACTOR);
log.debug(String.format("Fetched GMS groups with urns %s", existingGroups.keySet()));
@ -335,7 +336,7 @@ public class OidcCallbackLogic extends DefaultCallbackLogic<Result, PlayWebConte
// Now batch create all entities identified to create.
_entityClient.batchUpdate(groupsToCreate.stream().map(groupSnapshot ->
new Entity().setValue(Snapshot.create(groupSnapshot))
).collect(Collectors.toSet()), SYSTEM_PRINCIPAL);
).collect(Collectors.toSet()), SYSTEM_ACTOR);
log.debug(String.format("Successfully provisioned groups with urns %s", groupsToCreateUrns));
@ -349,7 +350,7 @@ public class OidcCallbackLogic extends DefaultCallbackLogic<Result, PlayWebConte
private void verifyPreProvisionedUser(CorpuserUrn urn) {
// Validate that the user exists in the system (there is more than just a key aspect for them, as of today).
try {
final Entity corpUser = _entityClient.get(urn, SYSTEM_PRINCIPAL);
final Entity corpUser = _entityClient.get(urn, SYSTEM_ACTOR);
log.debug(String.format("Fetched GMS user with urn %s", urn));

View File

@ -8,6 +8,7 @@ import com.linkedin.datahub.upgrade.impl.DefaultUpgradeStepResult;
import com.linkedin.datahub.upgrade.UpgradeContext;
import com.linkedin.datahub.upgrade.UpgradeStep;
import com.linkedin.datahub.upgrade.UpgradeStepResult;
import com.linkedin.metadata.Constants;
import com.linkedin.metadata.utils.PegasusUtils;
import com.linkedin.metadata.dao.utils.RecordUtils;
import com.linkedin.metadata.entity.EntityService;
@ -149,7 +150,7 @@ public class DataMigrationStep implements UpgradeStep {
browsePaths = BrowsePathUtils.buildBrowsePath(urn);
final AuditStamp browsePathsStamp = new AuditStamp();
browsePathsStamp.setActor(Urn.createFromString("urn:li:principal:system"));
browsePathsStamp.setActor(Urn.createFromString(Constants.SYSTEM_ACTOR));
browsePathsStamp.setTime(System.currentTimeMillis());
_entityService.ingestAspect(urn, BROWSE_PATHS_ASPECT_NAME, browsePaths, browsePathsStamp);

View File

@ -1,69 +0,0 @@
package com.linkedin.datahub.upgrade.nocode;
import com.linkedin.common.AuditStamp;
import com.linkedin.common.urn.DataPlatformUrn;
import com.linkedin.common.urn.Urn;
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.dataplatform.DataPlatformInfo;
import com.linkedin.metadata.utils.PegasusUtils;
import com.linkedin.metadata.entity.EntityService;
import com.linkedin.metadata.resources.dataplatform.utils.DataPlatformsUtil;
import java.net.URISyntaxException;
import java.time.Clock;
import java.util.Map;
import java.util.function.Function;
public class IngestDataPlatformsStep implements UpgradeStep {
private final EntityService _entityService;
public IngestDataPlatformsStep(final EntityService entityService) {
_entityService = entityService;
}
@Override
public String id() {
return "IngestDataPlatformsStep";
}
@Override
public int retryCount() {
return 2;
}
@Override
public Function<UpgradeContext, UpgradeStepResult> executable() {
return (context) -> {
context.report().addLine("Preparing to ingest DataPlatforms...");
Map<DataPlatformUrn, DataPlatformInfo> urnToInfo = DataPlatformsUtil.getDataPlatformInfoMap();
context.report().addLine(String.format("Found %s DataPlatforms", urnToInfo.keySet().size()));
for (final Map.Entry<DataPlatformUrn, DataPlatformInfo> entry : urnToInfo.entrySet()) {
AuditStamp auditStamp;
try {
auditStamp = new AuditStamp().setActor(Urn.createFromString("urn:li:principal:system")).setTime(
Clock.systemUTC().millis());
} catch (URISyntaxException e) {
throw new RuntimeException("Failed to create Actor Urn");
}
_entityService.ingestAspect(
entry.getKey(),
PegasusUtils.getAspectNameFromSchema(entry.getValue().schema()),
entry.getValue(),
auditStamp
);
}
context.report().addLine(String.format("Successfully ingested %s DataPlatforms.", urnToInfo.keySet().size()));
return new DefaultUpgradeStepResult(id(), UpgradeStepResult.Result.SUCCEEDED);
};
}
}

View File

@ -66,7 +66,6 @@ public class NoCodeUpgrade implements Upgrade {
steps.add(new GMSQualificationStep());
steps.add(new UpgradeQualificationStep(server));
steps.add(new CreateAspectTableStep(server));
steps.add(new IngestDataPlatformsStep(entityService));
steps.add(new DataMigrationStep(server, entityService, entityRegistry));
steps.add(new GMSEnableWriteModeStep(entityClient));
return steps;

View File

@ -61,7 +61,7 @@ public class UpgradeQualificationStep implements UpgradeStep {
boolean v2TableExists = AspectStorageValidationUtil.checkV2TableExists(server);
if (v2TableExists) {
context.report().addLine("-- V2 table exists");
long v2TableRowCount = AspectStorageValidationUtil.getV2RowCount(server);
long v2TableRowCount = AspectStorageValidationUtil.getV2NonSystemRowCount(server);
if (v2TableRowCount == 0) {
context.report().addLine("-- V2 table is empty");
return true;

View File

@ -17,12 +17,12 @@ insert into metadata_aspect_v2 (urn, aspect, version, metadata, createdon, creat
0,
'{"displayName":"Data Hub","active":true,"fullName":"Data Hub","email":"datahub@linkedin.com"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
), (
'urn:li:corpuser:datahub',
'corpUserEditableInfo',
0,
'{"skills":[],"teams":[],"pictureLink":"https://raw.githubusercontent.com/linkedin/datahub/master/datahub-web-react/src/images/default_avatar.png"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
);

View File

@ -23,14 +23,14 @@ INSERT INTO temp_metadata_aspect_v2 (urn, aspect, version, metadata, createdon,
0,
'{"displayName":"Data Hub","active":true,"fullName":"Data Hub","email":"datahub@linkedin.com"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
), (
'urn:li:corpuser:datahub',
'corpUserEditableInfo',
0,
'{"skills":[],"teams":[],"pictureLink":"https://raw.githubusercontent.com/linkedin/datahub/master/datahub-web-react/src/images/default_avatar.png"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
);
-- only add default records if metadata_aspect is empty
INSERT INTO metadata_aspect_v2

View File

@ -17,14 +17,14 @@ INSERT INTO metadata_aspect_v2 (urn, aspect, version, metadata, createdon, creat
0,
'{"displayName":"Data Hub","active":true,"fullName":"Data Hub","email":"datahub@linkedin.com"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
), (
'urn:li:corpuser:datahub',
'corpUserEditableInfo',
0,
'{"skills":[],"teams":[],"pictureLink":"https://raw.githubusercontent.com/linkedin/datahub/master/datahub-web-react/src/images/default_avatar.png"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
);
-- create metadata index table

View File

@ -19,14 +19,14 @@ INSERT INTO temp_metadata_aspect_v2 (urn, aspect, version, metadata, createdon,
0,
'{"displayName":"Data Hub","active":true,"fullName":"Data Hub","email":"datahub@linkedin.com"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
), (
'urn:li:corpuser:datahub',
'corpUserEditableInfo',
0,
'{"skills":[],"teams":[],"pictureLink":"https://raw.githubusercontent.com/linkedin/datahub/master/datahub-web-react/src/images/default_avatar.png"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
);
-- only add default records if metadata_aspect is empty
INSERT INTO metadata_aspect_v2

View File

@ -17,12 +17,12 @@ insert into metadata_aspect_v2 (urn, aspect, version, metadata, createdon, creat
0,
'{"displayName":"Data Hub","active":true,"fullName":"Data Hub","email":"datahub@linkedin.com"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
), (
'urn:li:corpuser:datahub',
'corpUserEditableInfo',
0,
'{"skills":[],"teams":[],"pictureLink":"https://raw.githubusercontent.com/linkedin/datahub/master/datahub-web-react/src/images/default_avatar.png"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
);

View File

@ -17,14 +17,14 @@ INSERT INTO metadata_aspect_v2 (urn, aspect, version, metadata, createdon, creat
0,
'{"displayName":"Data Hub","active":true,"fullName":"Data Hub","email":"datahub@linkedin.com"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
), (
'urn:li:corpuser:datahub',
'corpUserEditableInfo',
0,
'{"skills":[],"teams":[],"pictureLink":"https://raw.githubusercontent.com/linkedin/datahub/master/datahub-web-react/src/images/default_avatar.png"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
);
-- create metadata index table

View File

@ -19,14 +19,14 @@ insert into metadata_aspect (urn, aspect, version, metadata, createdon, createdb
0,
'{"displayName":"Data Hub","active":true,"fullName":"Data Hub","email":"datahub@linkedin.com"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
), (
'urn:li:corpuser:datahub',
'com.linkedin.identity.CorpUserEditableInfo',
0,
'{"skills":[],"teams":[],"pictureLink":"https://raw.githubusercontent.com/linkedin/datahub/master/datahub-web-react/src/images/default_avatar.png"}',
now(),
'urn:li:principal:datahub'
'urn:li:corpuser:__datahub_system'
);
-- create metadata index table

View File

@ -1,5 +1,6 @@
package com.linkedin.metadata.entity;
import com.linkedin.metadata.Constants;
import com.linkedin.metadata.entity.ebean.EbeanAspectV1;
import com.linkedin.metadata.entity.ebean.EbeanAspectV2;
import io.ebean.EbeanServer;
@ -7,6 +8,8 @@ import io.ebean.SqlQuery;
import io.ebean.SqlRow;
import java.util.List;
import static io.ebean.Expr.*;
public class AspectStorageValidationUtil {
@ -18,8 +21,11 @@ public class AspectStorageValidationUtil {
return server.find(EbeanAspectV1.class).findCount();
}
public static long getV2RowCount(EbeanServer server) {
return server.find(EbeanAspectV2.class).findCount();
/**
* Get the number of rows created not by the DataHub system actor (urn:li:corpuser:__datahub_system)
*/
public static long getV2NonSystemRowCount(EbeanServer server) {
return server.find(EbeanAspectV2.class).where(ne("createdby", Constants.SYSTEM_ACTOR)).findCount();
}
public static boolean checkV2TableExists(EbeanServer server) {

View File

@ -57,7 +57,7 @@ public class IngestDataPlatformsStep implements BootstrapStep {
RecordUtils.toRecordTemplate(DataPlatformInfo.class, dataPlatform.get("aspect").toString());
final AuditStamp aspectAuditStamp =
new AuditStamp().setActor(Urn.createFromString(Constants.UNKNOWN_ACTOR)).setTime(System.currentTimeMillis());
new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis());
_entityService.ingestAspect(urn, PLATFORM_ASPECT_NAME, info, aspectAuditStamp);
}

View File

@ -5,6 +5,7 @@ import com.fasterxml.jackson.databind.ObjectMapper;
import com.linkedin.common.AuditStamp;
import com.linkedin.common.urn.Urn;
import com.linkedin.data.template.RecordTemplate;
import com.linkedin.metadata.Constants;
import com.linkedin.metadata.boot.BootstrapStep;
import com.linkedin.events.metadata.ChangeType;
import com.linkedin.metadata.dao.utils.RecordUtils;
@ -92,7 +93,7 @@ public class IngestPoliciesStep implements BootstrapStep {
keyAspectProposal.setEntityUrn(urn);
_entityService.ingestProposal(keyAspectProposal,
new AuditStamp().setActor(Urn.createFromString("urn:li:corpuser:system")).setTime(System.currentTimeMillis()));
new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()));
final MetadataChangeProposal proposal = new MetadataChangeProposal();
proposal.setEntityUrn(urn);
@ -102,7 +103,7 @@ public class IngestPoliciesStep implements BootstrapStep {
proposal.setChangeType(ChangeType.UPSERT);
_entityService.ingestProposal(proposal,
new AuditStamp().setActor(Urn.createFromString("urn:li:corpuser:system")).setTime(System.currentTimeMillis()));
new AuditStamp().setActor(Urn.createFromString(Constants.SYSTEM_ACTOR)).setTime(System.currentTimeMillis()));
}
private boolean hasDefaultPolicies() throws URISyntaxException {

View File

@ -6,7 +6,7 @@ package com.linkedin.metadata;
public class Constants {
public static final String ACTOR_HEADER_NAME = "X-DataHub-Actor";
public static final String DATAHUB_ACTOR = "urn:li:corpuser:datahub"; // Super user.
public static final String SYSTEM_ACTOR = "urn:li:principal:datahub"; // DataHub internal service principal.
public static final String SYSTEM_ACTOR = "urn:li:corpuser:__datahub_system"; // DataHub internal service principal.
public static final String UNKNOWN_ACTOR = "urn:li:corpuser:UNKNOWN"; // Unknown principal.
public static final Long ASPECT_LATEST_VERSION = 0L;