From 67249ccf96014efe2eff226cb4487f995e2066f3 Mon Sep 17 00:00:00 2001
From: billettc <charles@cbillette.com>
Date: Tue, 7 Nov 2023 18:05:20 -0500
Subject: [PATCH] Fix missing AiTrainerPayments

---
 cmd/honey-tracker/main.go |  7 ++----
 data/db.go                | 45 +++++++++++++++++----------------------
 data/psql.go              | 20 +++++++++++++++++
 data/sinker.go            |  8 +++++--
 4 files changed, 48 insertions(+), 32 deletions(-)

diff --git a/cmd/honey-tracker/main.go b/cmd/honey-tracker/main.go
index 966d079..e219538 100644
--- a/cmd/honey-tracker/main.go
+++ b/cmd/honey-tracker/main.go
@@ -5,17 +5,14 @@ import (
 	"fmt"
 	"os"
 
-	"github.com/streamingfast/bstream"
-
 	"github.com/spf13/cobra"
+	"github.com/streamingfast/bstream"
 	"github.com/streamingfast/cli/sflags"
-
-	"go.uber.org/zap"
-
 	"github.com/streamingfast/honey-tracker/data"
 	"github.com/streamingfast/logging"
 	sink "github.com/streamingfast/substreams-sink"
 	"github.com/streamingfast/substreams/client"
+	"go.uber.org/zap"
 )
 
 var RootCmd = &cobra.Command{
diff --git a/data/db.go b/data/db.go
index 5bfd540..d27d930 100644
--- a/data/db.go
+++ b/data/db.go
@@ -1,27 +1,22 @@
 package data
 
-import (
-	pb "github.com/streamingfast/honey-tracker/data/pb/hivemapper/v1"
-	sink "github.com/streamingfast/substreams-sink"
-	pbsubstreams "github.com/streamingfast/substreams/pb/sf/substreams/v1"
-)
-
-type DB interface {
-	Init() error
-
-	BeginTransaction() error
-	CommitTransaction() error
-	RollbackTransaction() error
-
-	HandleClock(clock *pbsubstreams.Clock) (dbBlockID int64, err error)
-	HandleInitializedAccount(dbBlockID int64, initializedAccount []*pb.InitializedAccount) error
-	HandleRegularDriverPayments(dbBlockID int64, payments []*pb.RegularDriverPayment) error
-	HandleAiPayments(dbBlockID int64, payments []*pb.AiTrainerPayment) error
-	HandleSplitPayments(dbBlockID int64, splitPayments []*pb.TokenSplittingPayment) error
-	HandleNoneSplitPayments(dbBlockID int64, payments []*pb.NoSplitPayment) error
-	HandleTransfers(dbBlockID int64, transfers []*pb.Transfer) error
-	HandleMints(dbBlockID int64, mints []*pb.Mint) error
-	HandleBurns(dbBlockID int64, burns []*pb.Burn) error
-	StoreCursor(cursor *sink.Cursor) error
-	FetchCursor() (*sink.Cursor, error)
-}
+//type DB interface {
+//	Init() error
+//
+//	BeginTransaction() error
+//	CommitTransaction() error
+//	RollbackTransaction() error
+//
+//	HandleClock(clock *pbsubstreams.Clock) (dbBlockID int64, err error)
+//	HandleInitializedAccount(dbBlockID int64, initializedAccount []*pb.InitializedAccount) error
+//	HandleRegularDriverPayments(dbBlockID int64, payments []*pb.RegularDriverPayment) error
+//	HandleAiPayments(dbBlockID int64, payments []*pb.AiTrainerPayment) error
+//	HandleSplitPayments(dbBlockID int64, splitPayments []*pb.TokenSplittingPayment) error
+//	HandleNoneSplitPayments(dbBlockID int64, payments []*pb.NoSplitPayment) error
+//	HandleTransfers(dbBlockID int64, transfers []*pb.Transfer) error
+//	HandleMints(dbBlockID int64, mints []*pb.Mint) error
+//	HandleBurns(dbBlockID int64, burns []*pb.Burn) error
+//	StoreCursor(cursor *sink.Cursor) error
+//	FetchCursor() (*sink.Cursor, error)
+//	HandleAITrainerPayments(dbBlockID int64, payments []*pb.AiTrainerPayment) error
+//}
diff --git a/data/psql.go b/data/psql.go
index 89da6a2..d8eba22 100644
--- a/data/psql.go
+++ b/data/psql.go
@@ -324,6 +324,26 @@ func (p *Psql) HandleMints(dbBlockID int64, mints []*pb.Mint) error {
 	}
 	return nil
 }
+func (p *Psql) HandleAITrainerPayments(dbBlockID int64, payments []*pb.AiTrainerPayment) error {
+	for _, payment := range payments {
+		dbTransactionID, err := p.handleTransaction(dbBlockID, payment.Mint.TrxHash)
+		if err != nil {
+			return fmt.Errorf("inserting transaction: %w", err)
+		}
+
+		mintDbID, err := p.insertMint(dbTransactionID, payment.Mint)
+		if err != nil {
+			return fmt.Errorf("inserting mint: %w", err)
+		}
+
+		_, err = p.tx.Exec("INSERT INTO hivemapper.ai_payments (mint_id) VALUES ($1) RETURNING id", mintDbID)
+		if err != nil {
+			return fmt.Errorf("inserting payment: %w", err)
+		}
+	}
+
+	return nil
+}
 
 func (p *Psql) insertBurns(dbTransactionID int64, burn *pb.Burn) (dbMintID int64, err error) {
 	row := p.tx.QueryRow("INSERT INTO hivemapper.burns (transaction_id, from_address, amount) VALUES ($1, $2, $3) RETURNING id", dbTransactionID, burn.From, burn.Amount)
diff --git a/data/sinker.go b/data/sinker.go
index 475ef24..85705d2 100644
--- a/data/sinker.go
+++ b/data/sinker.go
@@ -16,12 +16,12 @@ import (
 type Sinker struct {
 	logger *zap.Logger
 	*sink.Sinker
-	db            DB
+	db            *Psql
 	lastClock     *v1.Clock
 	blockSecCount int64
 }
 
-func NewSinker(logger *zap.Logger, sink *sink.Sinker, db DB) *Sinker {
+func NewSinker(logger *zap.Logger, sink *sink.Sinker, db *Psql) *Sinker {
 	return &Sinker{
 		logger: logger,
 		Sinker: sink,
@@ -116,6 +116,10 @@ func (s *Sinker) HandleBlockScopedData(ctx context.Context, data *pbsubstreamsrp
 		return fmt.Errorf("handle payments: %w", err)
 	}
 
+	if err := s.db.HandleAITrainerPayments(dbBlockID, moduleOutput.AiTrainerPayments); err != nil {
+		return fmt.Errorf("handle AiTrainerPayments: %w", err)
+	}
+
 	if err := s.db.HandleSplitPayments(dbBlockID, moduleOutput.TokenSplittingPayments); err != nil {
 		return fmt.Errorf("handle split payments: %w", err)
 	}