pool-stats/database/db.go

package database

import (
	"fmt"
	"log"
	"pool-stats/models"
	"time"

	"github.com/ostafen/clover/v2"
	"github.com/ostafen/clover/v2/document"
	c "github.com/ostafen/clover/v2/query"
	badgerstore "github.com/ostafen/clover/v2/store/badger"
)

const (
	CollectionName = "shares"

	TopSharesCollectionName = "TopShares"
	TopSharesAmount         = 20

	TimeWindowHighShareCollectionName = "TimeWindowHighShareStat"
)

func InitDatabase(path string) (*clover.DB, error) {
	store, err := badgerstore.Open(path)
	if err != nil {
		return nil, fmt.Errorf("failed to open BadgerDB store: %v", err)
	}

	db, err := clover.OpenWithStore(store)
	if err != nil {
		return nil, fmt.Errorf("failed to open CloverDB: %v", err)
	}

	// Ensure collection exists
	hasCollection, err := db.HasCollection(CollectionName)
	if err != nil {
		return nil, fmt.Errorf("failed to check collection: %v", err)
	}

	if !hasCollection {
		if err := db.CreateCollection(CollectionName); err != nil {
			return nil, fmt.Errorf("failed to create collection: %v", err)
		}
		if err := db.CreateIndex(CollectionName, "CreateDate"); err != nil {
			return nil, fmt.Errorf("failed to create index: %v", err)
		}
	}

	// Init TopShares collection
	hasTopSharesCollection, err := db.HasCollection(TopSharesCollectionName)
	if err != nil {
		return nil, fmt.Errorf("failed to check TopShares collection: %v", err)
	}

	if !hasTopSharesCollection {
		if err := db.CreateCollection(TopSharesCollectionName); err != nil {
			return nil, fmt.Errorf("failed to create TopShares collection: %v", err)
		}

		if err := db.CreateIndex(TopSharesCollectionName, "CreateDate"); err != nil {
			return nil, fmt.Errorf("failed to create index for TopShares: %v", err)
		}

		if err := db.CreateIndex(TopSharesCollectionName, "SDiff"); err != nil {
			return nil, fmt.Errorf("failed to create index for TopShares SDiff: %v", err)
		}
	}

	// Init TimeWindowHighShareStat collection
	hasTimeWindowCollection, err := db.HasCollection(TimeWindowHighShareCollectionName)
	if err != nil {
		return nil, fmt.Errorf("failed to check TimeWindowHighShare collection: %v", err)
	}

	if !hasTimeWindowCollection {
		if err := db.CreateCollection(TimeWindowHighShareCollectionName); err != nil {
			return nil, fmt.Errorf("failed to create TimeWindowHighShare collection: %v", err)
		}
		if err := db.CreateIndex(TimeWindowHighShareCollectionName, "TimeWindowID"); err != nil {
			return nil, fmt.Errorf("failed to create index for TimeWindowHighShare: %v", err)
		}
	}

	return db, nil
}

func GetHighestSharesInRange(db *clover.DB, collection string, since time.Time, count int) ([]models.ShareLog, error) {
	// Convert `since` to the format in `createdate`
	lower := since.Unix()
	upper := time.Now().Unix()

	// Filter by timestamp range
	criteria := c.Field("CreateDate").GtEq(fmt.Sprint(lower)).
		And(c.Field("CreateDate").LtEq(fmt.Sprint(upper)))

	// Query sorted by "sdiff" descending, limit 1
	results, err := db.FindAll(c.NewQuery(collection).
		Where(criteria).
		Sort(c.SortOption{Field: "SDiff", Direction: -1}).
		Limit(count))

	if err != nil || len(results) == 0 {
		return nil, err
	}

	var shares []models.ShareLog
	for _, doc := range results {
		var s models.ShareLog
		if err := doc.Unmarshal(&s); err != nil {
			return nil, err
		}
		shares = append(shares, s)
	}

	return shares, nil
}

func PrintAllHashes(db *clover.DB) {
	docs, err := db.FindAll(c.NewQuery(CollectionName))
	if err != nil {
		log.Fatalf("Failed to read from collection: %v", err)
	}

	for _, doc := range docs {
		hash := doc.Get("Hash")
		fmt.Println(hash)
	}
}

func ListShares(db *clover.DB, offset int, count int) []models.ShareLog {
	results, err := db.FindAll(
		c.NewQuery(CollectionName).
			Sort(c.SortOption{Field: "CreateDate", Direction: -1}).
			Skip(offset).
			Limit(count),
	)
	if err != nil {
		log.Printf("failed to list shares: %v", err)
		return nil
	}

	shareLogs := make([]models.ShareLog, len(results))
	for idx, doc := range results {
		var shareLog models.ShareLog
		doc.Unmarshal(&shareLog)
		shareLogs[idx] = shareLog
	}

	return shareLogs
}

func ListTopShares(db *clover.DB) []models.ShareLog {
	results, err := db.FindAll(
		c.NewQuery(TopSharesCollectionName).
			Sort(c.SortOption{Field: "SDiff", Direction: -1}),
	)
	if err != nil {
		log.Printf("failed to list top shares: %v", err)
		return nil
	}

	topShares := make([]models.ShareLog, len(results))
	for idx, doc := range results {
		var shareLog models.ShareLog
		doc.Unmarshal(&shareLog)
		topShares[idx] = shareLog
	}

	return topShares
}

func ReplaceTopShares(db *clover.DB, shares []models.ShareLog) {
	db.Delete(c.NewQuery(TopSharesCollectionName))

	for _, share := range shares {
		doc := document.NewDocumentOf(&share)
		if _, err := db.InsertOne(TopSharesCollectionName, doc); err != nil {
			return
		}
	}

	log.Printf("Replaced TopShares with %d shares", len(shares))
}

func GetTimeWindowHighShares(db *clover.DB) []models.TimeWindowHighShare {
	results, err := db.FindAll(
		c.NewQuery(TimeWindowHighShareCollectionName).
			Sort(c.SortOption{Field: "TimeWindowID", Direction: 1}),
	)
	if err != nil {
		log.Printf("failed to list time window high shares: %v", err)
		return nil
	}

	timeWindowHighShares := make([]models.TimeWindowHighShare, len(results))
	for idx, doc := range results {
		var timeWindowHighShare models.TimeWindowHighShare
		doc.Unmarshal(&timeWindowHighShare)
		timeWindowHighShares[idx] = timeWindowHighShare
	}

	return timeWindowHighShares
}

func SetTimeWindowHighShare(db *clover.DB, share models.TimeWindowHighShare) error {
	doc := document.NewDocumentOf(&share)

	existingDoc, _ := db.FindFirst(c.NewQuery(TimeWindowHighShareCollectionName).
		Where(c.Field("TimeWindowID").Eq(share.TimeWindowID)))
	if existingDoc != nil {
		db.ReplaceById(TimeWindowHighShareCollectionName, existingDoc.ObjectId(), doc)
	} else {
		db.InsertOne(TimeWindowHighShareCollectionName, doc)
	}

	return nil
}