Admin CLI - Manual Rollback command (#436)

* feat: add rollback command in Admin CLI

* chore: update rollback yaml file

* chore: update rollback yml file

applications/admin-cli/src/main/java/com/bloxbean/cardano/yaci/store/admin/cli/db/DBCommands.java

@@ -3,8 +3,11 @@
 import com.bloxbean.cardano.yaci.store.admin.cli.Groups;
 import com.bloxbean.cardano.yaci.store.dbutils.index.model.IndexDefinition;
 import com.bloxbean.cardano.yaci.store.dbutils.index.model.TableIndex;
+import com.bloxbean.cardano.yaci.store.dbutils.index.model.TableRollbackAction;
 import com.bloxbean.cardano.yaci.store.dbutils.index.service.IndexService;
+import com.bloxbean.cardano.yaci.store.dbutils.index.service.RollbackService;
 import com.bloxbean.cardano.yaci.store.dbutils.index.util.IndexLoader;
+import com.bloxbean.cardano.yaci.store.dbutils.index.util.RollbackLoader;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.shell.command.annotation.Command;
@@ -20,8 +23,10 @@
 public class DBCommands {
     private static final String INDEX_FILE = "index.yml";
     private static final String EXTRA_INDEX_FILE = "extra-index.yml";
+    private static final String ROLLBACK_FILE = "rollback.yml";
 
     private final IndexService indexService;
+    private final RollbackService rollbackService;
 
     @Command(description = "Apply the default indexes required for read operations.")
     public void applyIndexes(@Option(longNames = "skip-extra-indexes", defaultValue = "false", description = "Skip additional optional indexes.") boolean skipExtraIndexes) {
@@ -40,6 +45,16 @@ public void applyExtraIndexes() {
         applyIndexes(EXTRA_INDEX_FILE);
     }
 
+    @Command(description = "Rollback data to a previous epoch")
+    public void rollbackData(@Option(longNames = "epoch", required = true, description = "Epoch to rollback to") int epoch,
+                         @Option(longNames = "event-publisher-id", defaultValue = "1000", description = "Event Publisher ID") long eventPublisherId) {
+        writeLn(info("Start to rollback data ..."));
+        if (isRollbackEpochValid(epoch)) {
+            verifyRollback(ROLLBACK_FILE);
+            applyRollback(ROLLBACK_FILE, epoch, eventPublisherId);
+        }
+    }
+
     private void applyIndexes(String indexFile) {
         IndexLoader indexLoader = new IndexLoader();
         List<IndexDefinition> indexDefinitionList = indexLoader.loadIndexes(indexFile);
@@ -84,4 +99,44 @@ private void verifyIndexes(String indexFile) {
         }
     }
 
+    private void applyRollback(String rollbackFile, int epoch, long eventPublisherId) {
+        RollbackLoader rollbackLoader = new RollbackLoader();
+
+        List<String> rollbackTableNames = rollbackLoader.loadRollbackTableNames(rollbackFile);
+
+        if (rollbackTableNames == null || rollbackTableNames.isEmpty()) {
+            log.warn("No table found to rollback");
+            writeLn(warn("No table found to rollback : {}", rollbackFile));
+            return;
+        }
+
+        var result = rollbackService.executeRollback(rollbackTableNames, epoch, eventPublisherId);
+
+        if (result.getSecond().equals(Boolean.FALSE)) {
+            log.warn(">> Failed to rollback data");
+            List<TableRollbackAction> failedTableRollbackActions = result.getFirst();
+
+            for (TableRollbackAction tableRollbackAction : failedTableRollbackActions) {
+                writeLn(warn("Failed to rollback table : %s, action : %s", tableRollbackAction.getTableName(), tableRollbackAction.getSql()));
+            }
+        } else {
+            writeLn(success("Data rollback is successful, data is rolled back to epoch : %d", epoch));
+        }
+    }
+
+    private void verifyRollback(String rollbackFile) {
+        RollbackLoader rollbackLoader = new RollbackLoader();
+
+        List<String> tableNames = rollbackLoader.loadRollbackTableNames(rollbackFile);
+        rollbackService.verifyRollbackActions(tableNames);
+    }
+
+    private boolean isRollbackEpochValid(int epoch) {
+        if (!rollbackService.isValidRollbackEpoch(epoch)) {
+            writeLn(warn("Epoch %d is not a valid rollback epoch. The rollback epoch must be " +
+                    "less than or equal to the current max epoch in the database and greater than 0.", epoch));
+            return false;
+        }
+        return true;
+    }
 }

components/dbutils/src/main/java/com/bloxbean/cardano/yaci/store/dbutils/index/model/RollbackBlock.java

@@ -0,0 +1,17 @@
+package com.bloxbean.cardano.yaci.store.dbutils.index.model;
+
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+
+@Data
+@Builder
+@AllArgsConstructor
+public class RollbackBlock {
+    private String hash;
+    private Long number;
+    private Long slot;
+    private Integer epoch;
+    private Integer epochSlot;
+    private Integer era;
+}

components/dbutils/src/main/java/com/bloxbean/cardano/yaci/store/dbutils/index/model/TableRollbackAction.java

@@ -0,0 +1,11 @@
+package com.bloxbean.cardano.yaci.store.dbutils.index.model;
+
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+
+@Getter
+@AllArgsConstructor
+public class TableRollbackAction {
+    private String tableName;
+    private String sql;
+}

components/dbutils/src/main/java/com/bloxbean/cardano/yaci/store/dbutils/index/service/RollbackService.java

@@ -0,0 +1,166 @@
+package com.bloxbean.cardano.yaci.store.dbutils.index.service;
+
+import com.bloxbean.cardano.yaci.store.dbutils.index.model.*;
+import com.bloxbean.cardano.yaci.store.dbutils.index.util.DatabaseUtils;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.data.util.Pair;
+import org.springframework.jdbc.core.namedparam.MapSqlParameterSource;
+import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
+import org.springframework.stereotype.Service;
+import org.springframework.transaction.annotation.Transactional;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@Service
+@RequiredArgsConstructor
+@Slf4j
+public class RollbackService {
+    private final NamedParameterJdbcTemplate jdbcTemplate;
+    private final DatabaseUtils databaseUtils;
+
+    @Transactional
+    public Pair<List<TableRollbackAction>, Boolean> executeRollback(List<String> tableNames, int epoch, long eventPublisherId) {
+        RollbackBlock rollbackBlock = getRollbackBlockByEpoch(epoch);
+
+        if (rollbackBlock == null) {
+            log.error("Failed to get rollback block for epoch: {}", epoch);
+            return Pair.of(new ArrayList<>(), false);
+        }
+
+        rollbackCursor(rollbackBlock, eventPublisherId);
+        rollbackAccountConfig(rollbackBlock);
+
+        var params = new MapSqlParameterSource();
+        params.addValue("epoch", rollbackBlock.getEpoch());
+        params.addValue("slot", rollbackBlock.getSlot());
+
+        List<TableRollbackAction> failedRollbackActions = new ArrayList<>();
+
+        // Execute DELETE statements for each table/condition
+        for (String tableName : tableNames) {
+            if (databaseUtils.tableExists(tableName)) {
+                String sql = buildDeleteSql(tableName, rollbackBlock.getEpoch(), rollbackBlock.getSlot());
+                log.info("Executing rollback on table '{}': {}", tableName, sql);
+                try {
+                    jdbcTemplate.update(sql, params);
+                } catch (Exception e) {
+                    log.error("Failed to execute rollback on table '{}': {}", tableName, e.getMessage());
+                    failedRollbackActions.add(new TableRollbackAction(tableName, sql));
+                }
+            }
+        }
+
+        boolean rollbackSuccess = failedRollbackActions.isEmpty();
+
+        return Pair.of(failedRollbackActions, rollbackSuccess);
+    }
+
+    private RollbackBlock getRollbackBlockByEpoch(int epoch) {
+        // TODO: Handling for cases where there is no 'block' table
+        String sql = "SELECT hash, slot, number, epoch_slot, era FROM block WHERE epoch = :epoch ORDER BY slot DESC LIMIT 1";
+
+        var params = new MapSqlParameterSource()
+                .addValue("epoch", epoch - 1);
+
+        return jdbcTemplate.queryForObject(sql, params, (rs, rowNum) ->
+                RollbackBlock.builder()
+                        .hash(rs.getString("hash"))
+                        .slot(rs.getLong("slot"))
+                        .epoch(epoch)
+                        .epochSlot(rs.getInt("epoch_slot"))
+                        .number(rs.getLong("number"))
+                        .era(Integer.valueOf(rs.getString("era")))
+                        .build()
+
+        );
+    }
+
+    private Integer getMaxEpoch() {
+        String sql = "SELECT MAX(epoch) FROM block";
+        return jdbcTemplate.getJdbcTemplate().queryForObject(sql, Integer.class);
+    }
+
+    public Pair<List<String>, List<String>> verifyRollbackActions(List<String> tableNames) {
+        List<String> tableExists = new ArrayList<>();
+        List<String> tableNotExists = new ArrayList<>();
+
+        for (String tableName: tableNames) {
+            if (databaseUtils.tableExists(tableName)) {
+                tableExists.add(tableName);
+            } else {
+                tableNotExists.add(tableName);
+            }
+        }
+
+        return Pair.of(tableExists, tableNotExists);
+    }
+
+    public boolean isValidRollbackEpoch(int epoch) {
+        Integer maxEpoch = getMaxEpoch();
+
+        if (maxEpoch == null) {
+            log.error("Failed to get max epoch from block table");
+            return false;
+        }
+
+        return epoch >= 1 && epoch <= maxEpoch;
+    }
+
+    private String buildDeleteSql(String table, int epoch, long slot) {
+        String deleteFilter = buildDeleteFilter(table, epoch, slot);
+        return "DELETE FROM " + table + " WHERE " + deleteFilter;
+    }
+
+    private String buildDeleteFilter(String tableName, int epoch, long slot) {
+        if (tableName.equals("epoch_stake") || tableName.equals("drep_dist") || tableName.equals("gov_action_proposal_status")) {
+            return "epoch >= " + epoch;
+        } else if (tableName.equals("tx_input")) {
+            return "spent_at_slot > " + slot;
+        }
+        else {
+            return "slot > " + slot;
+        }
+    }
+
+    private void rollbackCursor(RollbackBlock rollbackBlock, long eventPublisherId) {
+        String truncateCursor = "TRUNCATE TABLE cursor_";
+        log.info("Truncating cursor_ table: {}", truncateCursor);
+        jdbcTemplate.getJdbcTemplate().update(truncateCursor);
+
+        String insertCursor =
+                "INSERT INTO cursor_ (id, block_hash, slot, block_number, era) " +
+                        "VALUES (:id, :block_hash, :slot, :block_number, :era)";
+        log.info("Inserting into cursor_: {}", insertCursor);
+
+        var params = new MapSqlParameterSource();
+        params.addValue("id", eventPublisherId);
+        params.addValue("slot", rollbackBlock.getSlot());
+        params.addValue("block_number", rollbackBlock.getNumber());
+        params.addValue("block_hash", rollbackBlock.getHash());
+        params.addValue("era", rollbackBlock.getEra());
+
+        jdbcTemplate.update(insertCursor, params);
+    }
+
+    private void rollbackAccountConfig(RollbackBlock rollbackBlock) {
+        if (databaseUtils.tableExists("account_config")) {
+
+            String truncateAccCfg = "TRUNCATE TABLE account_config";
+            log.info("Truncating account_config: {}", truncateAccCfg);
+
+            jdbcTemplate.getJdbcTemplate().update(truncateAccCfg);
+
+            String insertAccCfg =
+                    "INSERT INTO account_config (config_id, status, slot, block, block_hash) " +
+                            "VALUES ('last_account_balance_processed_block', null, :slot, :block_number, :block_hash)";
+            log.info("Inserting into account_config: {}", insertAccCfg);
+
+            jdbcTemplate.update(insertAccCfg, new MapSqlParameterSource()
+                    .addValue("slot", rollbackBlock.getSlot())
+                    .addValue("block_number", rollbackBlock.getNumber())
+                    .addValue("block_hash", rollbackBlock.getHash()));
+        }
+    }
+}

components/dbutils/src/main/java/com/bloxbean/cardano/yaci/store/dbutils/index/util/RollbackLoader.java

@@ -0,0 +1,38 @@
+package com.bloxbean.cardano.yaci.store.dbutils.index.util;
+
+import org.yaml.snakeyaml.Yaml;
+
+import java.io.InputStream;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+public class RollbackLoader {
+
+    @SuppressWarnings("unchecked")
+    public List<String> loadRollbackTableNames(String yamlFilePath) {
+        Yaml yaml = new Yaml();
+
+        InputStream is = getClass().getClassLoader().getResourceAsStream(yamlFilePath);
+        if (is == null) {
+            throw new IllegalArgumentException("File not found: " + yamlFilePath);
+        }
+
+        Map<String, Object> root = yaml.load(is);
+        List<String> tables = new ArrayList<>();
+
+        Object tablesObj = root.get("tables");
+        if (tablesObj instanceof List) {
+            for (Object t : (List<Object>) tablesObj) {
+                if (t instanceof String) {
+                    String table = ((String) t).trim();
+                    if (!table.isEmpty()) {
+                        tables.add(table);
+                    }
+                }
+            }
+        }
+
+        return tables;
+    }
+}

components/dbutils/src/main/resources/rollback.yml

@@ -0,0 +1,54 @@
+tables:
+  # asset store
+  - asset
+  # block store
+  - block
+  # epoch store
+  - protocol_params_proposal
+  - epoch_param
+  - cost_model
+  # governance store
+  - gov_action_proposal
+  - delegation_vote
+  - drep_registration
+  - committee
+  - committee_deregistration
+  - committee_member
+  - committee_registration
+  - constitution
+  - voting_procedure
+  # metadata store
+  - transaction_metadata
+  # mir store
+  - mir
+  # script store
+  - script
+  - transaction_scripts
+  # staking store
+  - stake_registration
+  - delegation
+  - pool_registration
+  - pool_retirement
+  - pool
+  # transaction store
+  - transaction
+  - transaction_witness
+  - withdrawal
+  - invalid_transaction
+  # utxo store
+  - address_utxo
+  - tx_input
+  # other tables from account, adapot, governance-aggr module
+  - adapot
+  - address_balance
+  - address_tx_amount
+  - cost_model
+  - instant_reward
+  - reward
+  - reward_rest
+  - stake_address_balance
+  - adapot_jobs
+  - drep
+  - epoch_stake
+  - drep_dist
+  - gov_action_proposal_status