Feature: total_bytes

celenium-io · Dec 9, 2024 · 3697169 · 3697169
1 parent 7eed35c
commit 3697169
Show file tree

Hide file tree

Showing 9 changed files with 83 additions and 23 deletions.
diff --git a/internal/storage/postgres/core.go b/internal/storage/postgres/core.go
@@ -7,12 +7,14 @@ import (
 	"context"
 
 	models "github.com/celenium-io/astria-indexer/internal/storage"
+	"github.com/celenium-io/astria-indexer/internal/storage/postgres/migrations"
 	"github.com/dipdup-net/go-lib/config"
 	"github.com/dipdup-net/go-lib/database"
 	"github.com/dipdup-net/indexer-sdk/pkg/storage"
 	"github.com/dipdup-net/indexer-sdk/pkg/storage/postgres"
 	"github.com/pkg/errors"
 	"github.com/uptrace/bun"
+	"github.com/uptrace/bun/migrate"
 	"go.opentelemetry.io/otel/trace"
 )
 
@@ -135,19 +137,18 @@ func initDatabaseWithMigrations(ctx context.Context, conn *database.Bun) error {
 	return migrateDatabase(ctx, conn)
 }
 
-func migrateDatabase(_ context.Context, _ *database.Bun) error {
-	// migrator := migrate.NewMigrator(db.DB(), migrations.Migrations)
-	// if err := migrator.Init(ctx); err != nil {
-	// 	return err
-	// }
-	// if err := migrator.Lock(ctx); err != nil {
-	// 	return err
-	// }
-	// defer migrator.Unlock(ctx) //nolint:errcheck
-
-	// _, err := migrator.Migrate(ctx)
-	// return err
-	return nil
+func migrateDatabase(ctx context.Context, db *database.Bun) error {
+	migrator := migrate.NewMigrator(db.DB(), migrations.Migrations)
+	if err := migrator.Init(ctx); err != nil {
+		return err
+	}
+	if err := migrator.Lock(ctx); err != nil {
+		return err
+	}
+	defer migrator.Unlock(ctx) //nolint:errcheck
+
+	_, err := migrator.Migrate(ctx)
+	return err
 }
 
 func createHypertables(ctx context.Context, conn *database.Bun) error {

diff --git a/internal/storage/postgres/migrations/20241208_total_bytes.go b/internal/storage/postgres/migrations/20241208_total_bytes.go
@@ -0,0 +1,39 @@
+// SPDX-FileCopyrightText: 2024 PK Lab AG <[email protected]>
+// SPDX-License-Identifier: MIT
+
+package migrations
+
+import (
+	"context"
+
+	"github.com/celenium-io/astria-indexer/internal/storage"
+	"github.com/uptrace/bun"
+)
+
+func init() {
+	Migrations.MustRegister(upTotalBytes, downTotalBytes)
+}
+
+func upTotalBytes(ctx context.Context, db *bun.DB) error {
+	_, err := db.ExecContext(ctx, `ALTER TABLE IF EXISTS public.state ADD COLUMN IF NOT EXISTS total_bytes int8 NULL`)
+	if err != nil {
+		return err
+	}
+	_, err = db.ExecContext(ctx, `COMMENT ON COLUMN public.state.total_bytes IS 'Total rollup bytes'`)
+	if err != nil {
+		return err
+	}
+
+	var totalBytes int64
+	err = db.NewSelect().Model((*storage.Rollup)(nil)).ColumnExpr("sum(size)").Scan(ctx, &totalBytes)
+	if err != nil {
+		return err
+	}
+
+	_, err = db.NewUpdate().Set("total_bytes = ?", totalBytes).Model((*storage.State)(nil)).Where("name = 'dipdup_astria_indexer'").Exec(ctx)
+	return err
+}
+func downTotalBytes(ctx context.Context, db *bun.DB) error {
+	_, err := db.ExecContext(ctx, `ALTER TABLE IF EXISTS public.state DROP COLUMN IF NOT EXISTS total_bytes`)
+	return err
+}
diff --git a/internal/storage/postgres/migrations/migrations.go b/internal/storage/postgres/migrations/migrations.go
@@ -0,0 +1,10 @@
+// SPDX-FileCopyrightText: 2024 PK Lab AG <[email protected]>
+// SPDX-License-Identifier: MIT
+
+package migrations
+
+import (
+	"github.com/uptrace/bun/migrate"
+)
+
+var Migrations = migrate.NewMigrations()
diff --git a/internal/storage/state.go b/internal/storage/state.go
@@ -37,6 +37,7 @@ type State struct {
 	TotalValidators int             `bun:"total_validators"          comment:"Validators count"         json:"validators"`
 	TotalSupply     decimal.Decimal `bun:"total_supply,type:numeric" comment:"Total supply"             json:"supply"`
 	TotalBridges    int64           `bun:"total_bridges"             comment:"Count of bridges"         json:"bridges"`
+	TotalBytes      int64           `bun:"total_bytes"               comment:"Total rollup bytes"       json:"bytes"`
 }
 
 // TableName -

diff --git a/pkg/indexer/storage/rollup.go b/pkg/indexer/storage/rollup.go
@@ -14,19 +14,22 @@ func (module *Module) saveRollup(
 	tx storage.Transaction,
 	rollups map[string]*storage.Rollup,
 	rollupAddress map[string]*storage.RollupAddress,
-) (int64, error) {
+) (int64, int64, error) {
 	if len(rollups) == 0 {
-		return 0, nil
+		return 0, 0, nil
 	}
 
+	var totalSize int64
+
 	data := make([]*storage.Rollup, 0)
 	for _, value := range rollups {
 		data = append(data, value)
+		totalSize += value.Size
 	}
 
 	count, err := tx.SaveRollups(ctx, data...)
 	if err != nil {
-		return count, err
+		return count, totalSize, err
 	}
 
 	ra := make([]*storage.RollupAddress, 0)
@@ -36,8 +39,8 @@ func (module *Module) saveRollup(
 		ra = append(ra, value)
 	}
 	if err := tx.SaveRollupAddresses(ctx, ra...); err != nil {
-		return 0, err
+		return 0, 0, err
 	}
 
-	return count, nil
+	return count, totalSize, nil
 }
diff --git a/pkg/indexer/storage/state.go b/pkg/indexer/storage/state.go
@@ -8,7 +8,7 @@ import (
 	"github.com/celenium-io/astria-indexer/pkg/types"
 )
 
-func updateState(block *storage.Block, totalAccounts, totalRollups, totalBridges int64, state *storage.State) {
+func updateState(block *storage.Block, totalAccounts, totalRollups, totalBridges, totalBytes int64, state *storage.State) {
 	if types.Level(block.Id) <= state.LastHeight {
 		return
 	}
@@ -20,6 +20,7 @@ func updateState(block *storage.Block, totalAccounts, totalRollups, totalBridges
 	state.TotalAccounts += totalAccounts
 	state.TotalRollups += totalRollups
 	state.TotalBridges += totalBridges
+	state.TotalBytes += totalBytes
 	state.TotalSupply = state.TotalSupply.Add(block.Stats.SupplyChange)
 	state.ChainId = block.ChainId
 }
diff --git a/pkg/indexer/storage/state_test.go b/pkg/indexer/storage/state_test.go
@@ -16,6 +16,7 @@ func Test_updateState(t *testing.T) {
 		block         *storage.Block
 		totalAccounts int64
 		totalRollups  int64
+		totalBytes    int64
 		state         *storage.State
 	}
 
@@ -42,6 +43,7 @@ func Test_updateState(t *testing.T) {
 				},
 				totalAccounts: 10,
 				totalRollups:  11,
+				totalBytes:    12,
 				state: &storage.State{
 					Id:            1,
 					Name:          "test",
@@ -53,6 +55,7 @@ func Test_updateState(t *testing.T) {
 					TotalRollups:  12,
 					TotalSupply:   decimal.RequireFromString("1000"),
 					TotalBridges:  10,
+					TotalBytes:    12,
 				},
 			},
 			want: storage.State{
@@ -66,12 +69,13 @@ func Test_updateState(t *testing.T) {
 				TotalRollups:  23,
 				TotalSupply:   decimal.RequireFromString("1100"),
 				TotalBridges:  20,
+				TotalBytes:    24,
 			},
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			updateState(tt.args.block, tt.args.totalAccounts, tt.args.totalRollups, 10, tt.args.state)
+			updateState(tt.args.block, tt.args.totalAccounts, tt.args.totalRollups, 10, tt.args.totalBytes, tt.args.state)
 		})
 	}
 }
diff --git a/pkg/indexer/storage/storage.go b/pkg/indexer/storage/storage.go
@@ -176,7 +176,7 @@ func (module *Module) processBlockInTransaction(ctx context.Context, tx storage.
 		return state, err
 	}
 
-	totalRollups, err := module.saveRollup(ctx, tx, block.Rollups, block.RollupAddress)
+	totalRollups, totalBytes, err := module.saveRollup(ctx, tx, block.Rollups, block.RollupAddress)
 	if err != nil {
 		return state, err
 	}
@@ -211,7 +211,7 @@ func (module *Module) processBlockInTransaction(ctx context.Context, tx storage.
 		return state, err
 	}
 
-	updateState(block, totalAccounts, totalRollups, totalBridges, &state)
+	updateState(block, totalAccounts, totalRollups, totalBridges, totalBytes, &state)
 	if err := tx.Update(ctx, &state); err != nil {
 		return state, err
 	}

diff --git a/test/data/state.yml b/test/data/state.yml
@@ -8,4 +8,5 @@
   total_accounts: 6
   total_rollups: 1
   total_supply: 1000000000000000000000
-  total_bridges: 1
+  total_bridges: 1
+  total_bytes: 1232