Track unsettled usage

pkg/db/queries.sql

@@ -180,3 +180,22 @@ ON CONFLICT (address)
 	RETURNING
 		id;
 
+-- name: IncrementUnsettledUsage :exec
+INSERT INTO unsettled_usage(payer_id, originator_id, minutes_since_epoch, spend_picodollars)
+	VALUES (@payer_id, @originator_id, @minutes_since_epoch, @spend_picodollars)
+ON CONFLICT (payer_id, originator_id, minutes_since_epoch)
+	DO UPDATE SET
+		spend_picodollars = unsettled_usage.spend_picodollars + @spend_picodollars;
+
+-- name: GetPayerUnsettledUsage :one
+SELECT
+	SUM(spend_picodollars) AS total_spend_picodollars
+FROM
+	unsettled_usage
+WHERE
+	payer_id = @payer_id
+	AND (@minutes_since_epoch_gt::BIGINT = 0
+		OR minutes_since_epoch > @minutes_since_epoch_gt::BIGINT)
+	AND (@minutes_since_epoch_lt::BIGINT = 0
+		OR minutes_since_epoch < @minutes_since_epoch_lt::BIGINT);
+

pkg/db/tx.go

@@ -34,3 +34,39 @@ func RunInTx(
 	done = true
 	return tx.Commit()
 }
+
+func RunInTxWithResult[T any](
+	ctx context.Context,
+	db *sql.DB,
+	opts *sql.TxOptions,
+	fn func(ctx context.Context, txQueries *queries.Queries) (T, error),
+) (T, error) {
+	querier := queries.New(db)
+	tx, err := db.BeginTx(ctx, opts)
+	if err != nil {
+		var zero T
+		return zero, err
+	}
+
+	var done bool
+
+	defer func() {
+		if !done {
+			_ = tx.Rollback()
+		}
+	}()
+
+	result, err := fn(ctx, querier.WithTx(tx))
+	if err != nil {
+		var zero T
+		return zero, err
+	}
+
+	done = true
+	if err := tx.Commit(); err != nil {
+		var zero T
+		return zero, err
+	}
+
+	return result, nil
+}

pkg/db/unsettledUsage_test.go

@@ -0,0 +1,100 @@
+package db
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/xmtp/xmtpd/pkg/db/queries"
+	"github.com/xmtp/xmtpd/pkg/testutils"
+	"github.com/xmtp/xmtpd/pkg/utils"
+)
+
+func TestIncrementUnsettledUsage(t *testing.T) {
+	ctx := context.Background()
+	db, _, cleanup := testutils.NewDB(t, ctx)
+	defer cleanup()
+
+	querier := queries.New(db)
+	payerId := testutils.RandomInt32()
+	originatorId := testutils.RandomInt32()
+	minutesSinceEpoch := utils.MinutesSinceEpochNow()
+
+	require.NoError(t, querier.IncrementUnsettledUsage(ctx, queries.IncrementUnsettledUsageParams{
+		PayerID:           payerId,
+		OriginatorID:      originatorId,
+		MinutesSinceEpoch: minutesSinceEpoch,
+		SpendPicodollars:  100,
+	}))
+
+	unsettledUsage, err := querier.GetPayerUnsettledUsage(
+		ctx,
+		queries.GetPayerUnsettledUsageParams{
+			PayerID: payerId,
+		},
+	)
+	require.NoError(t, err)
+	require.Equal(t, unsettledUsage, int64(100))
+
+	require.NoError(t, querier.IncrementUnsettledUsage(ctx, queries.IncrementUnsettledUsageParams{
+		PayerID:           payerId,
+		OriginatorID:      originatorId,
+		MinutesSinceEpoch: minutesSinceEpoch,
+		SpendPicodollars:  100,
+	}))
+
+	unsettledUsage, err = querier.GetPayerUnsettledUsage(
+		ctx,
+		queries.GetPayerUnsettledUsageParams{
+			PayerID: payerId,
+		},
+	)
+	require.NoError(t, err)
+	require.Equal(t, unsettledUsage, int64(200))
+}
+
+func TestGetUnsettledUsage(t *testing.T) {
+	ctx := context.Background()
+	db, _, cleanup := testutils.NewDB(t, ctx)
+	defer cleanup()
+
+	querier := queries.New(db)
+	payerId := testutils.RandomInt32()
+	originatorId := testutils.RandomInt32()
+
+	addUsage := func(minutesSinceEpoch int32, spendPicodollars int64) {
+		require.NoError(
+			t,
+			querier.IncrementUnsettledUsage(ctx, queries.IncrementUnsettledUsageParams{
+				PayerID:           payerId,
+				OriginatorID:      originatorId,
+				MinutesSinceEpoch: minutesSinceEpoch,
+				SpendPicodollars:  spendPicodollars,
+			}),
+		)
+	}
+
+	addUsage(1, 100)
+	addUsage(2, 200)
+	addUsage(3, 300)
+
+	unsettledUsage, err := querier.GetPayerUnsettledUsage(
+		ctx,
+		queries.GetPayerUnsettledUsageParams{
+			PayerID:             payerId,
+			MinutesSinceEpochGt: 2,
+		},
+	)
+	require.NoError(t, err)
+	require.Equal(t, unsettledUsage, int64(300))
+
+	unsettledUsage, err = querier.GetPayerUnsettledUsage(
+		ctx,
+		queries.GetPayerUnsettledUsageParams{
+			PayerID:             payerId,
+			MinutesSinceEpochGt: 1,
+		},
+	)
+	require.NoError(t, err)
+	require.Equal(t, unsettledUsage, int64(500))
+}

pkg/migrations/00007_unsettled-usage.down.sql

@@ -0,0 +1,4 @@
+DROP TABLE IF EXISTS unsettled_usage;
+
+DROP INDEX IF EXISTS idx_unsettled_usage_payer_id;
+

pkg/migrations/00007_unsettled-usage.up.sql

@@ -0,0 +1,9 @@
+CREATE TABLE unsettled_usage(
+	payer_id INTEGER NOT NULL,
+	originator_id INTEGER NOT NULL,
+	minutes_since_epoch INTEGER NOT NULL,
+	spend_picodollars BIGINT NOT NULL,
+	PRIMARY KEY (payer_id, originator_id, minutes_since_epoch)
+);
+
+CREATE INDEX idx_unsettled_usage_payer_id ON unsettled_usage(payer_id);

pkg/testutils/random.go

@@ -67,3 +67,7 @@ func RandomBlockHash() common.Hash {
 	bytes := RandomBytes(32)
 	return common.BytesToHash(bytes)
 }
+
+func RandomInt32() int32 {
+	return rand.Int31()
+}

pkg/utils/time.go

@@ -0,0 +1,13 @@
+package utils
+
+import "time"
+
+func MinutesSinceEpoch(timestamp time.Time) int32 {
+	durationSinceEpoch := timestamp.Sub(time.Unix(0, 0))
+
+	return int32(durationSinceEpoch.Minutes())
+}
+
+func MinutesSinceEpochNow() int32 {
+	return MinutesSinceEpoch(time.Now())
+}