diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/UpgradeQualificationStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/UpgradeQualificationStep.java index 56ecda604f443a..0912cbb4d73010 100644 --- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/UpgradeQualificationStep.java +++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/UpgradeQualificationStep.java @@ -32,40 +32,49 @@ public Function executable() { if (context.parsedArgs().containsKey(NoCodeUpgrade.FORCE_UPGRADE_ARG_NAME)) { context.report().addLine("Forced upgrade detected. Proceeding with upgrade..."); - return new DefaultUpgradeStepResult( - id(), - UpgradeStepResult.Result.SUCCEEDED); + return new DefaultUpgradeStepResult(id(), UpgradeStepResult.Result.SUCCEEDED); } try { - if (isQualified(_server)) { + if (isQualified(_server, context)) { // Qualified. context.report().addLine("Found qualified upgrade candidate. Proceeding with upgrade..."); - return new DefaultUpgradeStepResult( - id(), - UpgradeStepResult.Result.SUCCEEDED); + return new DefaultUpgradeStepResult(id(), UpgradeStepResult.Result.SUCCEEDED); } // Unqualified (Table already exists) context.report().addLine("Failed to qualify upgrade candidate. Aborting the upgrade..."); - return new DefaultUpgradeStepResult( - id(), - UpgradeStepResult.Result.SUCCEEDED, - UpgradeStepResult.Action.ABORT); + return new DefaultUpgradeStepResult(id(), UpgradeStepResult.Result.SUCCEEDED, UpgradeStepResult.Action.ABORT); } catch (Exception e) { context.report().addLine(String.format("Failed to check if metadata_aspect_v2 table exists: %s", e.toString())); - return new DefaultUpgradeStepResult( - id(), - UpgradeStepResult.Result.FAILED); + return new DefaultUpgradeStepResult(id(), UpgradeStepResult.Result.FAILED); } }; } - private boolean isQualified(EbeanServer server) { - - return AspectStorageValidationUtil.checkV1TableExists(server) - && ( - !AspectStorageValidationUtil.checkV2TableExists(server) - || AspectStorageValidationUtil.getV2RowCount(server) == 0 - ); + // Check whether the upgrade is needed + private boolean isQualified(EbeanServer server, UpgradeContext context) { + boolean v1TableExists = AspectStorageValidationUtil.checkV1TableExists(server); + if (v1TableExists) { + context.report().addLine("-- V1 table exists"); + long v1TableRowCount = AspectStorageValidationUtil.getV1RowCount(server); + context.report().addLine(String.format("-- V1 table has %d rows", v1TableRowCount)); + boolean v2TableExists = AspectStorageValidationUtil.checkV2TableExists(server); + if (v2TableExists) { + context.report().addLine("-- V2 table exists"); + long v2TableRowCount = AspectStorageValidationUtil.getV2RowCount(server); + if (v2TableRowCount == 0) { + context.report().addLine("-- V2 table is empty"); + return true; + } + context.report().addLine(String.format("-- V2 table has %d rows", v2TableRowCount)); + context.report().addLine("-- Since V2 table has records, we will not proceed with the upgrade. "); + context.report().addLine("-- If V2 table has significantly less rows, consider running the forced upgrade. "); + return false; + } + context.report().addLine("-- V2 table does not exist"); + return true; + } + context.report().addLine("-- V1 table does not exist"); + return false; } } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectStorageValidationUtil.java b/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectStorageValidationUtil.java index df3845b47fa2e6..ef034b723bb990 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectStorageValidationUtil.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/entity/AspectStorageValidationUtil.java @@ -1,5 +1,6 @@ package com.linkedin.metadata.entity; +import com.linkedin.metadata.entity.ebean.EbeanAspectV1; import com.linkedin.metadata.entity.ebean.EbeanAspectV2; import io.ebean.EbeanServer; import io.ebean.SqlQuery; @@ -11,6 +12,10 @@ private AspectStorageValidationUtil() { } + public static long getV1RowCount(EbeanServer server) { + return server.find(EbeanAspectV1.class).findCount(); + } + public static long getV2RowCount(EbeanServer server) { return server.find(EbeanAspectV2.class).findCount(); }