validationMigration publish to new kafka

Author: nischitpra

cmd/migrate_valid.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"context"
+	"fmt"
 	"math/big"
 	"os"
 	"strconv"
@@ -11,6 +12,7 @@ import (
 	config "github.com/thirdweb-dev/indexer/configs"
 	"github.com/thirdweb-dev/indexer/internal/common"
 	"github.com/thirdweb-dev/indexer/internal/orchestrator"
+	"github.com/thirdweb-dev/indexer/internal/publisher/newkafka"
 	"github.com/thirdweb-dev/indexer/internal/rpc"
 	"github.com/thirdweb-dev/indexer/internal/storage"
 )
@@ -36,6 +38,9 @@ func RunValidationMigration(cmd *cobra.Command, args []string) {
 	migrator := NewMigrator()
 	defer migrator.Close()
 
+	// get absolute start and end block for the migration. eg, 0-10M or 10M-20M
+	absStartBlock, absEndBlock := migrator.getAbsStartAndEndBlock()
+
 	rangeStartBlock, rangeEndBlock := migrator.DetermineMigrationBoundaries()
 
 	log.Info().Msgf("Migrating blocks from %s to %s (both ends inclusive)", rangeStartBlock.String(), rangeEndBlock.String())
@@ -75,9 +80,14 @@ func RunValidationMigration(cmd *cobra.Command, args []string) {
 			blocksToInsert = append(blocksToInsert, blockData)
 		}
 
-		err := migrator.targetConn.InsertBlockData(blocksToInsert)
+		err := migrator.newkafka.PublishBlockData(blocksToInsert)
 		if err != nil {
-			log.Fatal().Err(err).Msg("Failed to insert blocks to target storage")
+			log.Fatal().Err(err).Msg("Failed to publish block data")
+		}
+
+		err = migrator.UpdateMigratedBlock(absStartBlock, absEndBlock, blocksToInsert)
+		if err != nil {
+			log.Fatal().Err(err).Msg("Failed to update migrated block range")
 		}
 
 		currentBlock = new(big.Int).Add(endBlock, big.NewInt(1))
@@ -94,6 +104,8 @@ type Migrator struct {
 	targetConn         *storage.ClickHouseConnector
 	migrationBatchSize int
 	rpcBatchSize       int
+	newkafka           *newkafka.Publisher
+	psql               *storage.PostgresConnector
 }
 
 func NewMigrator() *Migrator {
@@ -138,21 +150,97 @@ func NewMigrator() *Migrator {
 		log.Fatal().Err(err).Msg("Failed to initialize target storage")
 	}
 
+	// publish to new kafka stream i.e new clickhouse database
+	newpublisher := newkafka.GetInstance()
+
+	// psql cursor for new kafka
+	psql, err := storage.NewPostgresConnector(config.Cfg.Storage.Main.Postgres)
+	if err != nil {
+		log.Fatal().Err(err).Msg("Failed to initialize psql cursor")
+	}
+
+	// Create migrated_block_ranges table if it doesn't exist
+	createMigratedBlockRangesTable(psql)
+
 	return &Migrator{
 		migrationBatchSize: batchSize,
 		rpcBatchSize:       rpcBatchSize,
 		rpcClient:          rpcClient,
 		storage:            s,
 		validator:          validator,
 		targetConn:         targetConn,
+		newkafka:           newpublisher,
+		psql:               psql,
+	}
+}
+
+// createMigratedBlockRangesTable creates the migrated_block_ranges table if it doesn't exist
+func createMigratedBlockRangesTable(psql *storage.PostgresConnector) {
+	createTableSQL := `
+		CREATE TABLE IF NOT EXISTS migrated_block_ranges (
+			chain_id BIGINT NOT NULL,
+			block_number BIGINT NOT NULL,
+			created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+			updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+		) WITH (fillfactor = 80, autovacuum_vacuum_scale_factor = 0.1, autovacuum_analyze_scale_factor = 0.05)
+	`
+
+	// Execute the CREATE TABLE statement
+	_, err := psql.ExecRaw(createTableSQL)
+	if err != nil {
+		log.Warn().Err(err).Msg("Failed to create migrated_block_ranges table")
+	}
+
+	// Create index if it doesn't exist
+	createIndexSQL := `
+		CREATE INDEX IF NOT EXISTS idx_migrated_block_ranges_chain_block 
+		ON migrated_block_ranges(chain_id, block_number DESC)
+	`
+	_, err = psql.ExecRaw(createIndexSQL)
+	if err != nil {
+		log.Warn().Err(err).Msg("Failed to create index on migrated_block_ranges table")
+	}
+
+	// Create trigger if it doesn't exist
+	createTriggerSQL := `
+		CREATE TRIGGER update_migrated_block_ranges_updated_at 
+		BEFORE UPDATE ON migrated_block_ranges 
+		FOR EACH ROW EXECUTE FUNCTION update_updated_at_column()
+	`
+	_, err = psql.ExecRaw(createTriggerSQL)
+	if err != nil {
+		log.Warn().Err(err).Msg("Failed to create trigger on migrated_block_ranges table")
 	}
 }
 
 func (m *Migrator) Close() {
 	m.rpcClient.Close()
+	m.newkafka.Close()
+	m.psql.Close()
 }
 
 func (m *Migrator) DetermineMigrationBoundaries() (*big.Int, *big.Int) {
+	startBlock, endBlock := m.getAbsStartAndEndBlock()
+
+	latestMigratedBlock, err := m.GetMaxBlockNumberInRange(startBlock, endBlock)
+	if err != nil {
+		log.Fatal().Err(err).Msg("Failed to get latest block from target storage")
+	}
+	log.Info().Msgf("Latest block in target storage: %d", latestMigratedBlock)
+
+	if latestMigratedBlock.Cmp(endBlock) >= 0 {
+		log.Fatal().Msgf("Full range is already migrated")
+	}
+
+	// if configured start block is less than or equal to already migrated and migrated block is not 0, start from last migrated + 1
+	if startBlock.Cmp(latestMigratedBlock) <= 0 && latestMigratedBlock.Sign() > 0 {
+		startBlock = new(big.Int).Add(latestMigratedBlock, big.NewInt(1))
+	}
+
+	return startBlock, endBlock
+}
+
+func (m *Migrator) getAbsStartAndEndBlock() (*big.Int, *big.Int) {
 	// get latest block from main storage
 	latestBlockStored, err := m.storage.MainStorage.GetMaxBlockNumber(m.rpcClient.GetChainID())
 	if err != nil {
@@ -187,22 +275,102 @@ func (m *Migrator) DetermineMigrationBoundaries() (*big.Int, *big.Int) {
 		startBlock = configuredStartBlock
 	}
 
-	latestMigratedBlock, err := m.targetConn.GetMaxBlockNumberInRange(m.rpcClient.GetChainID(), startBlock, endBlock)
+	return startBlock, endBlock
+}
+
+func (m *Migrator) UpdateMigratedBlock(startBlock *big.Int, endBlock *big.Int, blockData []common.BlockData) error {
+	if len(blockData) == 0 {
+		return nil
+	}
+
+	maxBlockNumber := big.NewInt(0)
+	for _, block := range blockData {
+		if block.Block.Number.Cmp(maxBlockNumber) > 0 {
+			maxBlockNumber = block.Block.Number
+		}
+	}
+
+	chainID := blockData[0].Block.ChainId
+	err := m.upsertMigratedBlockRange(chainID, maxBlockNumber, startBlock, endBlock)
 	if err != nil {
-		log.Fatal().Err(err).Msg("Failed to get latest block from target storage")
+		return fmt.Errorf("failed to update migrated block range: %w", err)
 	}
-	log.Info().Msgf("Latest block in target storage: %d", latestMigratedBlock)
+	return nil
+}
 
-	if latestMigratedBlock.Cmp(endBlock) >= 0 {
-		log.Fatal().Msgf("Full range is already migrated")
+func (m *Migrator) GetMaxBlockNumberInRange(startBlock *big.Int, endBlock *big.Int) (*big.Int, error) {
+	// Get chain ID from RPC client
+	chainID := m.rpcClient.GetChainID()
+
+	// Get the maximum end_block for the given chain_id
+	maxBlock, err := m.getMaxMigratedBlock(chainID, startBlock, endBlock)
+	if err != nil {
+		log.Warn().Err(err).Msg("Failed to get last migrated block, returning start block")
+		return startBlock, err
 	}
 
-	// if configured start block is less than or equal to already migrated and migrated block is not 0, start from last migrated + 1
-	if startBlock.Cmp(latestMigratedBlock) <= 0 && latestMigratedBlock.Sign() > 0 {
-		startBlock = new(big.Int).Add(latestMigratedBlock, big.NewInt(1))
+	// Return maxBlock + 1 as the next block to migrate
+	return new(big.Int).Add(maxBlock, big.NewInt(1)), nil
+}
+
+// upsertMigratedBlockRange upserts a row for the given chain_id and block range
+func (m *Migrator) upsertMigratedBlockRange(chainID, blockNumber, startBlock, endBlock *big.Int) error {
+	// First, try to update existing rows that overlap with this range
+	updateSQL := `
+		UPDATE migrated_block_ranges 
+		SET block_number = $1, updated_at = CURRENT_TIMESTAMP
+		WHERE chain_id = $2 AND block_number >= $3 AND block_number <= $4
+	`
+
+	result, err := m.psql.ExecRaw(updateSQL, blockNumber.String(), chainID.String(), startBlock.String(), endBlock.String())
+	if err != nil {
+		return fmt.Errorf("failed to update migrated block range for chain %s, range %s-%s", chainID.String(), startBlock.String(), endBlock.String())
 	}
 
-	return startBlock, endBlock
+	// Check if any rows were updated
+	rowsAffected, err := result.RowsAffected()
+	if err != nil {
+		log.Warn().Err(err).Msgf("Failed to get rows affected for chain %s, range %s-%s", chainID.String(), startBlock.String(), endBlock.String())
+		return fmt.Errorf("failed to get rows affected for chain %s, range %s-%s", chainID.String(), startBlock.String(), endBlock.String())
+	}
+
+	// If no rows were updated, insert a new row
+	if rowsAffected == 0 {
+		insertSQL := `
+			INSERT INTO migrated_block_ranges (chain_id, block_number, created_at, updated_at)
+			VALUES ($1, $2, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
+		`
+
+		_, err := m.psql.ExecRaw(insertSQL, chainID.String(), blockNumber.String())
+		if err != nil {
+			return fmt.Errorf("failed to insert migrated block range for chain %s, range %s-%s", chainID.String(), startBlock.String(), endBlock.String())
+		}
+	}
+	return nil
+}
+
+// getMaxMigratedBlock gets the maximum block number within the given range for the given chain_id
+func (m *Migrator) getMaxMigratedBlock(chainID, startBlock, endBlock *big.Int) (*big.Int, error) {
+	querySQL := `
+		SELECT COALESCE(MAX(block_number), 0) as max_block
+		FROM migrated_block_ranges
+		WHERE chain_id = $1 
+		AND block_number >= $2 
+		AND block_number <= $3
+	`
+
+	var maxBlockStr string
+	err := m.psql.QueryRowRaw(querySQL, chainID.String(), startBlock.String(), endBlock.String()).Scan(&maxBlockStr)
+	if err != nil {
+		return nil, fmt.Errorf("failed to query migrated block ranges: %w", err)
+	}
+
+	maxBlock, ok := new(big.Int).SetString(maxBlockStr, 10)
+	if !ok {
+		return nil, fmt.Errorf("failed to parse block number: %s", maxBlockStr)
+	}
+
+	return maxBlock, nil
 }
 
 func (m *Migrator) FetchBlocksFromRPC(blockNumbers []*big.Int) ([]common.BlockData, error) {