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

datahub-frontend/app/auth/sso/oidc/OidcCallbackLogic.java

@@ -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));
 

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java

@@ -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);

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/IngestDataPlatformsStep.java

@@ -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);
-    };
-  }
-}

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/NoCodeUpgrade.java

@@ -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;

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/UpgradeQualificationStep.java

@@ -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;

docker/mariadb/init.sql

@@ -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'
 );

docker/mysql-setup/init.sql

@@ -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

docker/mysql/init.sql

@@ -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

docker/postgres-setup/init.sql

@@ -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

docker/postgres/init.sql

@@ -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'
 );

docker/quickstart/mysql/init.sql

@@ -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

metadata-ingestion/tests/integration/mysql/setup/setup.sql

@@ -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

metadata-io/src/main/java/com/linkedin/metadata/entity/AspectStorageValidationUtil.java

@@ -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) {

metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestDataPlatformsStep.java

@@ -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);
     }

metadata-service/factories/src/main/java/com/linkedin/metadata/boot/steps/IngestPoliciesStep.java

@@ -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 {

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

@@ -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;