Rerun conditional migrations

docs/database/README.md

@@ -230,3 +230,11 @@ Following are the prerequisites and steps for migrating V1 data to V2.
 6. Stop the [Importer](/docs/importer/README.md) process.
 7. Run the [migration.sh](/hedera-mirror-importer/src/main/resources/db/scripts/v2/migration.sh) script.
 8. Update the mirror node configuration to point to the new Citus DB and start it.
+
+## Repeatable migration
+The following migrations would be rerun automatically when the first account balance file is parsed:
+* InitializeEntityBalanceMigration
+* TokenAccountBalanceMigration
+
+If there is an error during the rerun, or for some reason the migration does not run, it can be manually rerun using the steps provided at:
+[Importer-Initialize Entity Balance](/docs/importer/README.md#Initialize Entity Balance)

hedera-mirror-importer/src/main/java/com/hedera/mirror/importer/migration/ErrataMigration.java

@@ -30,6 +30,7 @@
 import com.hedera.mirror.importer.parser.record.entity.EntityProperties;
 import com.hedera.mirror.importer.parser.record.entity.EntityRecordItemListener;
 import com.hedera.mirror.importer.reader.ValidatedDataInputStream;
+import com.hedera.mirror.importer.repository.AccountBalanceFileRepository;
 import com.hedera.mirror.importer.repository.TokenTransferRepository;
 import com.hedera.mirror.importer.repository.TransactionRepository;
 import com.hederahashgraph.api.proto.java.Transaction;
@@ -67,6 +68,7 @@ public class ErrataMigration extends RepeatableMigration implements BalanceStrea
     @Value("classpath:errata/mainnet/balance-offsets.txt")
     private final Resource balanceOffsets;
 
+    private final AccountBalanceFileRepository accountBalanceFileRepository;
     private final EntityRecordItemListener entityRecordItemListener;
     private final EntityProperties entityProperties;
     private final NamedParameterJdbcOperations jdbcOperations;
@@ -81,6 +83,7 @@ public class ErrataMigration extends RepeatableMigration implements BalanceStrea
     @SuppressWarnings("java:S107")
     public ErrataMigration(
             Resource balanceOffsets,
+            AccountBalanceFileRepository accountBalanceFileRepository,
             EntityRecordItemListener entityRecordItemListener,
             EntityProperties entityProperties,
             NamedParameterJdbcOperations jdbcOperations,
@@ -91,6 +94,7 @@ public ErrataMigration(
             TransactionRepository transactionRepository) {
         super(mirrorProperties.getMigration());
         this.balanceOffsets = balanceOffsets;
+        this.accountBalanceFileRepository = accountBalanceFileRepository;
         this.entityRecordItemListener = entityRecordItemListener;
         this.entityProperties = entityProperties;
         this.jdbcOperations = jdbcOperations;
@@ -122,6 +126,7 @@ public void onEnd(AccountBalanceFile accountBalanceFile) {
             if (shouldApplyFixedTimeOffset(consensusTimestamp)) {
                 accountBalanceFile.setTimeOffset(ACCOUNT_BALANCE_FILE_FIXED_TIME_OFFSET);
             }
+            accountBalanceFileRepository.save(accountBalanceFile);
         }
     }
 

hedera-mirror-importer/src/main/java/com/hedera/mirror/importer/migration/InitializeEntityBalanceMigration.java

@@ -18,13 +18,17 @@
 
 import com.google.common.base.Stopwatch;
 import com.hedera.mirror.importer.MirrorProperties;
+import com.hedera.mirror.importer.repository.AccountBalanceFileRepository;
+import com.hedera.mirror.importer.repository.RecordFileRepository;
 import jakarta.inject.Named;
 import org.flywaydb.core.api.MigrationVersion;
 import org.springframework.context.annotation.Lazy;
+import org.springframework.core.annotation.Order;
 import org.springframework.jdbc.core.JdbcOperations;
 
 @Named
-public class InitializeEntityBalanceMigration extends RepeatableMigration {
+@Order
+public class InitializeEntityBalanceMigration extends TimeSensitiveBalanceMigration {
 
     private static final String INITIALIZE_ENTITY_BALANCE_SQL =
             """
@@ -66,8 +70,12 @@ on conflict (id) do update
     private final JdbcOperations jdbcOperations;
 
     @Lazy
-    public InitializeEntityBalanceMigration(JdbcOperations jdbcOperations, MirrorProperties mirrorProperties) {
-        super(mirrorProperties.getMigration());
+    public InitializeEntityBalanceMigration(
+            JdbcOperations jdbcOperations,
+            MirrorProperties mirrorProperties,
+            AccountBalanceFileRepository accountBalanceFileRepository,
+            RecordFileRepository recordFileRepository) {
+        super(mirrorProperties.getMigration(), accountBalanceFileRepository, recordFileRepository);
         this.jdbcOperations = jdbcOperations;
     }
 

hedera-mirror-importer/src/main/java/com/hedera/mirror/importer/migration/SyntheticCryptoTransferApprovalMigration.java

@@ -17,32 +17,43 @@
 package com.hedera.mirror.importer.migration;
 
 import com.hedera.mirror.common.domain.entity.EntityId;
+import com.hedera.mirror.common.domain.transaction.RecordFile;
 import com.hedera.mirror.importer.MirrorProperties;
 import com.hedera.mirror.importer.db.DBProperties;
+import com.hedera.mirror.importer.exception.ImporterException;
+import com.hedera.mirror.importer.parser.record.RecordStreamFileListener;
+import com.hedera.mirror.importer.repository.RecordFileRepository;
 import com.hederahashgraph.api.proto.java.Key;
 import jakarta.inject.Named;
+import java.io.IOException;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
 import lombok.Data;
 import lombok.Getter;
 import org.flywaydb.core.api.MigrationVersion;
 import org.springframework.context.annotation.Lazy;
+import org.springframework.data.util.Version;
 import org.springframework.jdbc.core.DataClassRowMapper;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
 import org.springframework.transaction.support.TransactionOperations;
 
 @Named
-public class SyntheticCryptoTransferApprovalMigration extends AsyncJavaMigration<Long> {
+public class SyntheticCryptoTransferApprovalMigration extends AsyncJavaMigration<Long>
+        implements RecordStreamFileListener {
 
+    static final Version HAPI_VERSION_0_38_0 = new Version(0, 38, 0);
+    private final AtomicBoolean executed = new AtomicBoolean(false);
     // The contract id of the first synthetic transfer that could have exhibited this problem
     private static final long GRANDFATHERED_ID = 2119900L;
     // The created timestamp of the grandfathered id contract
     static final long LOWER_BOUND_TIMESTAMP = 1680284879342064922L;
-    // This problem was fixed by services release 0.38.10, this is last timestamp before that release
+    // This problem was fixed by services release 0.38.6,protobuf release 0.38.10, this is last timestamp before that
+    // release
     static final long UPPER_BOUND_TIMESTAMP = 1686243920981874002L;
     private static final long TIMESTAMP_INCREMENT =
             Duration.ofDays(1).toNanos(); // 1 day in nanoseconds which will yield 69 async iterations
@@ -141,6 +152,30 @@
             update token_transfer set is_approval = true where consensus_timestamp = :consensus_timestamp and token_id = :token_id and account_id = :sender
             """;
 
+    private final RecordFileRepository recordFileRepository;
+
+    @Override
+    public void onEnd(RecordFile streamFile) throws ImporterException {
+        if (streamFile == null) {
+            return;
+        }
+
+        // The services version 0.38.0 has the fixes this migration solves.
+        try {
+            if (streamFile.getHapiVersion().isGreaterThanOrEqualTo(HAPI_VERSION_0_38_0)
+                    && executed.compareAndSet(false, true)) {
+                var latestFile = recordFileRepository.findLatestBefore(streamFile.getConsensusStart());
+                if (latestFile
+                        .filter(f -> f.getHapiVersion().isLessThan(HAPI_VERSION_0_38_0))
+                        .isPresent()) {
+                    doMigrate();
+                }
+            }
+        } catch (IOException e) {
+            log.error("Error executing the migration again after consensus_timestamp {}", streamFile.getConsensusEnd());
+        }
+    }
+
     private enum TRANSFER_TYPE {
         CRYPTO_TRANSFER,
         NFT_TRANSFER,
@@ -159,10 +194,12 @@
     @Lazy
     public SyntheticCryptoTransferApprovalMigration(
             DBProperties dbProperties,
+            RecordFileRepository recordFileRepository,
             MirrorProperties mirrorProperties,
             NamedParameterJdbcTemplate transferJdbcTemplate,
             TransactionOperations transactionOperations) {
         super(mirrorProperties.getMigration(), transferJdbcTemplate, dbProperties.getSchema());
+        this.recordFileRepository = recordFileRepository;
         this.mirrorProperties = mirrorProperties;
         this.transferJdbcTemplate = transferJdbcTemplate;
         this.transactionOperations = transactionOperations;

hedera-mirror-importer/src/main/java/com/hedera/mirror/importer/migration/SyntheticNftAllowanceOwnerMigration.java

@@ -17,15 +17,25 @@
 package com.hedera.mirror.importer.migration;
 
 import com.google.common.base.Stopwatch;
+import com.hedera.mirror.common.domain.transaction.RecordFile;
 import com.hedera.mirror.importer.MirrorProperties;
 import com.hedera.mirror.importer.config.Owner;
+import com.hedera.mirror.importer.exception.ImporterException;
+import com.hedera.mirror.importer.parser.record.RecordStreamFileListener;
+import com.hedera.mirror.importer.repository.RecordFileRepository;
 import jakarta.inject.Named;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.flywaydb.core.api.MigrationVersion;
 import org.springframework.context.annotation.Lazy;
+import org.springframework.data.util.Version;
 import org.springframework.jdbc.core.JdbcTemplate;
 
 @Named
-public class SyntheticNftAllowanceOwnerMigration extends RepeatableMigration {
+public class SyntheticNftAllowanceOwnerMigration extends RepeatableMigration implements RecordStreamFileListener {
+
+    static final Version HAPI_VERSION_0_37_0 = new Version(0, 37, 0);
+    private final AtomicBoolean executed = new AtomicBoolean(false);
+    private final RecordFileRepository recordFileRepository;
 
     private static final String UPDATE_NFT_ALLOWANCE_OWNER_SQL =
             """
@@ -107,9 +117,13 @@ insert into nft_allowance (approved_for_all, owner, payer_account_id, spender, t
     private final JdbcTemplate jdbcTemplate;
 
     @Lazy
-    public SyntheticNftAllowanceOwnerMigration(@Owner JdbcTemplate jdbcTemplate, MirrorProperties mirrorProperties) {
+    public SyntheticNftAllowanceOwnerMigration(
+            @Owner JdbcTemplate jdbcTemplate,
+            MirrorProperties mirrorProperties,
+            RecordFileRepository recordFileRepository) {
         super(mirrorProperties.getMigration());
         this.jdbcTemplate = jdbcTemplate;
+        this.recordFileRepository = recordFileRepository;
     }
 
     @Override
@@ -129,4 +143,22 @@ protected void doMigrate() {
         jdbcTemplate.execute(UPDATE_NFT_ALLOWANCE_OWNER_SQL);
         log.info("Updated nft allowance owners in {}", stopwatch);
     }
+
+    @Override
+    public void onEnd(RecordFile streamFile) throws ImporterException {
+        if (streamFile == null) {
+            return;
+        }
+
+        // The services version 0.37.0 has the fixes this migration solves.
+        if (streamFile.getHapiVersion().isGreaterThanOrEqualTo(HAPI_VERSION_0_37_0)
+                && executed.compareAndSet(false, true)) {
+            var latestFile = recordFileRepository.findLatestBefore(streamFile.getConsensusStart());
+            if (latestFile
+                    .filter(f -> f.getHapiVersion().isLessThan(HAPI_VERSION_0_37_0))
+                    .isPresent()) {
+                doMigrate();
+            }
+        }
+    }
 }

hedera-mirror-importer/src/main/java/com/hedera/mirror/importer/migration/SyntheticTokenAllowanceOwnerMigration.java

@@ -17,15 +17,25 @@
 package com.hedera.mirror.importer.migration;
 
 import com.google.common.base.Stopwatch;
+import com.hedera.mirror.common.domain.transaction.RecordFile;
 import com.hedera.mirror.importer.MirrorProperties;
 import com.hedera.mirror.importer.config.Owner;
+import com.hedera.mirror.importer.exception.ImporterException;
+import com.hedera.mirror.importer.parser.record.RecordStreamFileListener;
+import com.hedera.mirror.importer.repository.RecordFileRepository;
 import jakarta.inject.Named;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.flywaydb.core.api.MigrationVersion;
 import org.springframework.context.annotation.Lazy;
+import org.springframework.data.util.Version;
 import org.springframework.jdbc.core.JdbcTemplate;
 
 @Named
-public class SyntheticTokenAllowanceOwnerMigration extends RepeatableMigration {
+public class SyntheticTokenAllowanceOwnerMigration extends RepeatableMigration implements RecordStreamFileListener {
+
+    static final Version HAPI_VERSION_0_37_0 = new Version(0, 37, 0);
+    private final AtomicBoolean executed = new AtomicBoolean(false);
+    private final RecordFileRepository recordFileRepository;
 
     private static final String UPDATE_TOKEN_ALLOWANCE_OWNER_SQL =
             """
@@ -110,9 +120,13 @@ insert into token_allowance (amount, amount_granted, owner, payer_account_id, sp
     private final JdbcTemplate jdbcTemplate;
 
     @Lazy
-    public SyntheticTokenAllowanceOwnerMigration(@Owner JdbcTemplate jdbcTemplate, MirrorProperties mirrorProperties) {
+    public SyntheticTokenAllowanceOwnerMigration(
+            @Owner JdbcTemplate jdbcTemplate,
+            MirrorProperties mirrorProperties,
+            RecordFileRepository recordFileRepository) {
         super(mirrorProperties.getMigration());
         this.jdbcTemplate = jdbcTemplate;
+        this.recordFileRepository = recordFileRepository;
     }
 
     @Override
@@ -132,4 +146,22 @@ protected void doMigrate() {
         jdbcTemplate.execute(UPDATE_TOKEN_ALLOWANCE_OWNER_SQL);
         log.info("Updated token allowance owners in {}", stopwatch);
     }
+
+    @Override
+    public void onEnd(RecordFile streamFile) throws ImporterException {
+        if (streamFile == null) {
+            return;
+        }
+
+        // The services version 0.38.0 has the fixes this migration solves.
+        if (streamFile.getHapiVersion().isGreaterThanOrEqualTo(HAPI_VERSION_0_37_0)
+                && executed.compareAndSet(false, true)) {
+            var latestFile = recordFileRepository.findLatestBefore(streamFile.getConsensusStart());
+            if (latestFile
+                    .filter(f -> f.getHapiVersion().isLessThan(HAPI_VERSION_0_37_0))
+                    .isPresent()) {
+                doMigrate();
+            }
+        }
+    }
 }

hedera-mirror-importer/src/main/java/com/hedera/mirror/importer/migration/TimeSensitiveBalanceMigration.java

@@ -0,0 +1,92 @@
+/*
+ * Copyright (C) 2023 Hedera Hashgraph, LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.hedera.mirror.importer.migration;
+
+import com.hedera.mirror.common.domain.balance.AccountBalanceFile;
+import com.hedera.mirror.common.domain.transaction.RecordFile;
+import com.hedera.mirror.importer.exception.ImporterException;
+import com.hedera.mirror.importer.parser.balance.BalanceStreamFileListener;
+import com.hedera.mirror.importer.repository.AccountBalanceFileRepository;
+import com.hedera.mirror.importer.repository.RecordFileRepository;
+import java.io.IOException;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+import org.springframework.transaction.support.TransactionSynchronization;
+import org.springframework.transaction.support.TransactionSynchronizationManager;
+
+abstract class TimeSensitiveBalanceMigration extends RepeatableMigration
+        implements BalanceStreamFileListener, TransactionSynchronization {
+
+    private final AccountBalanceFileRepository accountBalanceFileRepository;
+    private final RecordFileRepository recordFileRepository;
+    private final AtomicLong firstConsensusTimestamp = new AtomicLong(0L);
+    private final AtomicBoolean succeeded = new AtomicBoolean(false);
+
+    protected TimeSensitiveBalanceMigration(
+            Map<String, MigrationProperties> migrationPropertiesMap,
+            AccountBalanceFileRepository accountBalanceFileRepository,
+            RecordFileRepository recordFileRepository) {
+        super(migrationPropertiesMap);
+        this.accountBalanceFileRepository = accountBalanceFileRepository;
+        this.recordFileRepository = recordFileRepository;
+    }
+
+    @Override
+    public void onEnd(AccountBalanceFile accountBalanceFile) throws ImporterException {
+        try {
+            if (firstConsensusTimestamp.get() == -1 || succeeded.get()) {
+                return;
+            }
+
+            if (firstConsensusTimestamp.get() == 0) {
+                // Check if this is the first account balance file after importer startup
+                // Set current file timestamp to firstConsensusTimestamp.
+                if (accountBalanceFileRepository
+                        .findLatestBefore(accountBalanceFile.getConsensusTimestamp())
+                        .isEmpty()) {
+                    firstConsensusTimestamp.set(accountBalanceFile.getConsensusTimestamp());
+                } else {
+                    // Set firstConsensusTimestamp to -1 to add an early return in case of existing account balance
+                    // files.
+                    firstConsensusTimestamp.set(-1);
+                    return;
+                }
+            }
+
+            // Check if at-least one recordFile after the account balance file has been parsed,the migration will then
+            // update rows
+            if (recordFileRepository
+                    .findLatest()
+                    .map(RecordFile::getConsensusEnd)
+                    .filter(timestamp -> timestamp >= firstConsensusTimestamp.get())
+                    .isPresent()) {
+                TransactionSynchronizationManager.registerSynchronization(this);
+                doMigrate();
+            }
+        } catch (IOException e) {
+            log.error(
+                    "Error executing the migration again after consensus_timestamp {}",
+                    accountBalanceFile.getConsensusTimestamp());
+        }
+    }
+
+    @Override
+    public void afterCommit() {
+        succeeded.set(true);
+    }
+}