pool-stats/cmd/migrate/main.go

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/ostafen/clover/v2"
	"github.com/ostafen/clover/v2/document"
	"github.com/ostafen/clover/v2/query"
	badgerstore "github.com/ostafen/clover/v2/store/badger"

	"pool-stats/database"
	"pool-stats/models"
)

const batchSize = 2000

var shareCopyColumns = []string{
	"work_info_id", "client_id", "enonce1", "nonce2", "nonce", "ntime",
	"diff", "sdiff", "hash", "result", "errn",
	"create_date", "create_ts", "create_by", "create_code", "create_inet",
	"worker_name", "username", "address", "agent",
}

func main() {
	badgerPath := flag.String("BadgerPath", "badgerdb", "Path to the BadgerDB/Clover directory")
	databaseDSN := flag.String("DatabaseDSN",
		"postgres://admin:password@127.0.0.1:5432/pkstats?sslmode=disable",
		"PostgreSQL connection string")
	force := flag.Bool("Force", false, "Truncate existing Postgres tables before migrating")
	flag.Parse()

	ctx := context.Background()

	log.Printf("Opening BadgerDB at %s...", *badgerPath)
	store, err := badgerstore.Open(*badgerPath)
	if err != nil {
		log.Fatalf("failed to open BadgerDB: %v", err)
	}
	cloverDB, err := clover.OpenWithStore(store)
	if err != nil {
		log.Fatalf("failed to open CloverDB: %v", err)
	}
	defer cloverDB.Close()

	log.Printf("Connecting to PostgreSQL...")
	pool, err := pgxpool.New(ctx, *databaseDSN)
	if err != nil {
		log.Fatalf("failed to connect to PostgreSQL: %v", err)
	}
	defer pool.Close()

	if err := database.ApplySchema(pool); err != nil {
		log.Fatalf("failed to apply schema: %v", err)
	}

	count, err := database.CountShares(pool)
	if err != nil {
		log.Fatalf("failed to count shares: %v", err)
	}
	if count > 0 {
		if !*force {
			log.Fatalf("refusing to migrate: shares table already has %d rows (use -Force to truncate and restart)", count)
		}
		log.Printf("Truncating existing tables (-Force)...")
		if err := truncateAll(ctx, pool); err != nil {
			log.Fatalf("failed to truncate: %v", err)
		}
	}

	shareCount, err := migrateShares(ctx, cloverDB, pool)
	if err != nil {
		log.Fatalf("failed to migrate shares: %v", err)
	}
	log.Printf("Migrated %d shares", shareCount)

	topCount, err := migrateTopShares(ctx, cloverDB, pool)
	if err != nil {
		log.Fatalf("failed to migrate top shares: %v", err)
	}
	log.Printf("Migrated %d top shares", topCount)

	twCount, err := migrateTimeWindowHighShares(ctx, cloverDB, pool)
	if err != nil {
		log.Fatalf("failed to migrate time window high shares: %v", err)
	}
	log.Printf("Migrated %d time window high shares", twCount)

	dailyCount, err := migrateDailyStats(ctx, cloverDB, pool)
	if err != nil {
		log.Fatalf("failed to migrate daily stats: %v", err)
	}
	log.Printf("Migrated %d daily stats", dailyCount)

	log.Println("Migration complete")
}

func truncateAll(ctx context.Context, pool *pgxpool.Pool) error {
	_, err := pool.Exec(ctx, `
		TRUNCATE TABLE shares, top_shares, time_window_high_shares, daily_stats
		RESTART IDENTITY`)
	return err
}

func migrateShares(ctx context.Context, cloverDB *clover.DB, pool *pgxpool.Pool) (int, error) {
	expected, err := cloverDB.Count(query.NewQuery(database.CollectionName))
	if err != nil {
		return 0, fmt.Errorf("count shares: %w", err)
	}
	log.Printf("Streaming %d shares from Badger...", expected)

	total := 0
	batch := make([][]any, 0, batchSize)
	lastLog := time.Now()
	var walkErr error

	flush := func() error {
		if len(batch) == 0 {
			return nil
		}
		_, err := pool.CopyFrom(
			ctx,
			pgx.Identifier{"shares"},
			shareCopyColumns,
			pgx.CopyFromRows(batch),
		)
		if err != nil {
			return err
		}
		total += len(batch)
		batch = batch[:0]
		if time.Since(lastLog) >= 2*time.Second || total == expected {
			log.Printf("  shares progress: %d / %d", total, expected)
			lastLog = time.Now()
		}
		return nil
	}

	err = cloverDB.ForEach(query.NewQuery(database.CollectionName), func(doc *document.Document) bool {
		var share models.ShareLog
		if uerr := doc.Unmarshal(&share); uerr != nil {
			walkErr = fmt.Errorf("unmarshal share: %w", uerr)
			return false
		}
		vals, verr := database.ShareRowValues(share)
		if verr != nil {
			walkErr = verr
			return false
		}
		batch = append(batch, vals)
		if len(batch) >= batchSize {
			if ferr := flush(); ferr != nil {
				walkErr = ferr
				return false
			}
		}
		return true
	})
	if err != nil {
		return total, err
	}
	if walkErr != nil {
		return total, walkErr
	}
	if err := flush(); err != nil {
		return total, err
	}
	return total, nil
}

func migrateTopShares(ctx context.Context, cloverDB *clover.DB, pool *pgxpool.Pool) (int, error) {
	total := 0
	batch := make([][]any, 0, batchSize)
	var walkErr error

	err := cloverDB.ForEach(query.NewQuery(database.TopSharesCollectionName), func(doc *document.Document) bool {
		var share models.ShareLog
		if uerr := doc.Unmarshal(&share); uerr != nil {
			walkErr = fmt.Errorf("unmarshal top share: %w", uerr)
			return false
		}
		vals, verr := database.ShareRowValues(share)
		if verr != nil {
			walkErr = verr
			return false
		}
		batch = append(batch, vals)
		total++
		return true
	})
	if err != nil {
		return 0, err
	}
	if walkErr != nil {
		return 0, walkErr
	}
	if len(batch) == 0 {
		return 0, nil
	}

	_, err = pool.CopyFrom(
		ctx,
		pgx.Identifier{"top_shares"},
		shareCopyColumns,
		pgx.CopyFromRows(batch),
	)
	return total, err
}

func migrateTimeWindowHighShares(ctx context.Context, cloverDB *clover.DB, pool *pgxpool.Pool) (int, error) {
	count := 0
	var walkErr error
	err := cloverDB.ForEach(query.NewQuery(database.TimeWindowHighShareCollectionName), func(doc *document.Document) bool {
		var tw models.TimeWindowHighShare
		if uerr := doc.Unmarshal(&tw); uerr != nil {
			walkErr = fmt.Errorf("unmarshal time window high share: %w", uerr)
			return false
		}
		if serr := database.SetTimeWindowHighShare(pool, tw); serr != nil {
			walkErr = serr
			return false
		}
		count++
		return true
	})
	if err != nil {
		return count, err
	}
	return count, walkErr
}

func migrateDailyStats(ctx context.Context, cloverDB *clover.DB, pool *pgxpool.Pool) (int, error) {
	count := 0
	var walkErr error
	err := cloverDB.ForEach(query.NewQuery(database.DailyStatsCollectionName), func(doc *document.Document) bool {
		var stats models.DailyStats
		if uerr := doc.Unmarshal(&stats); uerr != nil {
			walkErr = fmt.Errorf("unmarshal daily stats: %w", uerr)
			return false
		}

		topShareJSON, jerr := json.Marshal(stats.TopShare)
		if jerr != nil {
			walkErr = jerr
			return false
		}
		workersJSON, jerr := json.Marshal(stats.Workers)
		if jerr != nil {
			walkErr = jerr
			return false
		}

		var expires any
		if expiresAt := doc.ExpiresAt(); expiresAt != nil {
			expires = *expiresAt
		}

		dateStr := stats.Date
		if dateStr == "" {
			dateStr = time.Unix(0, 0).UTC().Format(time.DateOnly)
		}

		_, ierr := pool.Exec(ctx, `
			INSERT INTO daily_stats (date, share_count, top_share, pool_hashrate, workers, expires_at)
			VALUES ($1::date, $2, $3, $4, $5, $6)
			ON CONFLICT (date) DO UPDATE SET
				share_count = EXCLUDED.share_count,
				top_share = EXCLUDED.top_share,
				pool_hashrate = EXCLUDED.pool_hashrate,
				workers = EXCLUDED.workers,
				expires_at = EXCLUDED.expires_at`,
			dateStr, stats.ShareCount, topShareJSON, stats.PoolHashrate, workersJSON, expires)
		if ierr != nil {
			walkErr = fmt.Errorf("insert daily stats for %s: %w", dateStr, ierr)
			return false
		}
		count++
		return true
	})
	if err != nil {
		return count, err
	}
	return count, walkErr
}