Batch UPDATE в PostgreSQL

package postgres

import (
	"context"

	"github.com/doutivity/research-online-postgres-go/internal/storage/postgres/dbs"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Database - repository
type Database struct {
	pool    *pgxpool.Pool
	queries *dbs.Queries
}

// NewDatabase - constructor
func NewDatabase(pool *pgxpool.Pool) *Database {
	var queries = dbs.New(pool)

	return &Database{
		pool:    pool,
		queries: queries,
	}
}

// Connection - getter
func (r *Database) Connection() *pgxpool.Pool {
	return r.pool
}

// Queries - getter
func (r *Database) Queries() *dbs.Queries {
	return r.queries
}

// WithTransaction - start transaction
func (r *Database) WithTransaction(ctx context.Context, fn func(queries *dbs.Queries) error) error {
	return withTransaction(ctx, r.pool, r.queries, fn)
}

func withTransaction(
	ctx context.Context,
	db *pgxpool.Pool,
	queries *dbs.Queries,
	fn func(queries *dbs.Queries) error,
) (err error) {
	tx, err := db.BeginTx(ctx, pgx.TxOptions{})
	if err != nil {
		return
	}

	defer func() {
		if p := recover(); p != nil {
			// a panic occurred, rollback and repanic
			tx.Rollback(ctx)

			panic(p)
		} else if err != nil {
			// something went wrong, rollback
			tx.Rollback(ctx)
		} else {
			// all good, commit
			err = tx.Commit(ctx)
		}
	}()

	err = fn(queries.WithTx(tx))

	return err
}

package dbs

import (
	"context"

	"github.com/jackc/pgx/v5/pgtype"
)

// ...

type UserOnlineUpsertParams struct {
	UserID int64
	Online pgtype.Timestamp
}

func (q *Queries) UserOnlineUpsert(ctx context.Context, arg UserOnlineUpsertParams) error {
	_, err := q.db.Exec(ctx, userOnlineUpsert, arg.UserID, arg.Online)
	return err
}

// ...

type UserOnlineUpdateParams struct {
	Online pgtype.Timestamp
	UserID int64
}

func (q *Queries) UserOnlineUpdate(ctx context.Context, arg UserOnlineUpdateParams) error {
	_, err := q.db.Exec(ctx, userOnlineUpdate, arg.Online, arg.UserID)
	return err
}

// ...

type UserOnlineUnnestUpsertParams struct {
	UserIds []int64
	Onlines []pgtype.Timestamp
}

func (q *Queries) UserOnlineUnnestUpsert(ctx context.Context, arg UserOnlineUnnestUpsertParams) error {
	_, err := q.db.Exec(ctx, userOnlineUnnestUpsert, arg.UserIds, arg.Onlines)
	return err
}

// ...

type UserOnlineUnnestUpdateParams struct {
	UserIds []int64
	Onlines []pgtype.Timestamp
}

func (q *Queries) UserOnlineUnnestUpdate(ctx context.Context, arg UserOnlineUnnestUpdateParams) error {
	_, err := q.db.Exec(ctx, userOnlineUnnestUpdate, arg.UserIds, arg.Onlines)
	return err
}

package main

import (
	"context"
)

type OnlineStorage interface {
	BatchStore(ctx context.Context, pairs []UserOnlinePair) error
}

package main

type UserOnlinePair struct {
	UserID    int64
	Timestamp int64
}

package main

import (
	"context"
	"testing"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/stretchr/testify/require"
)

func testOnlineStorage(
	t *testing.T,
	storage OnlineStorage,
) {
	t.Helper()

	ctx := context.Background()

	pool, err := pgxpool.New(ctx, dataSourceName)
	require.NoError(t, err)
	defer pool.Close()

	var (
		pair1v1 = UserOnlinePair{
			UserID:    1,
			Timestamp: 2e3,
		}
		pair2v1 = UserOnlinePair{
			UserID:    2,
			Timestamp: 3e4,
		}
		pair3v1 = UserOnlinePair{
			UserID:    3,
			Timestamp: 4e5,
		}
		pair4v1 = UserOnlinePair{
			UserID:    4,
			Timestamp: 5e6,
		}
		pair5v1 = UserOnlinePair{
			UserID:    5,
			Timestamp: 6e7,
		}

		pair2v2 = UserOnlinePair{
			UserID:    pair2v1.UserID,
			Timestamp: pair2v1.Timestamp + 10001,
		}
		pair3v2 = UserOnlinePair{
			UserID:    pair3v1.UserID,
			Timestamp: pair3v1.Timestamp + 10002,
		}
	)

	truncateOnline(t, ctx, pool)
	insertOnline(t, ctx, pool, []UserOnlinePair{
		pair1v1,
		pair2v1,
		pair3v1,
		pair4v1,
		pair5v1,
	})

	err = storage.BatchStore(ctx, []UserOnlinePair{
		pair2v2,
		pair3v2,
	})
	require.NoError(t, err)

	expectedOnline(t, ctx, pool, []UserOnlinePair{
		pair1v1,
		pair2v2,
		pair3v2,
		pair4v1,
		pair5v1,
	})
}

package main

import (
	"context"
	"time"

	"github.com/doutivity/research-online-postgres-go/internal/storage/postgres"
	"github.com/doutivity/research-online-postgres-go/internal/storage/postgres/dbs"

	"github.com/jackc/pgx/v5/pgtype"
)

type TxLoopUpsertOnlineStorage struct {
	db *postgres.Database
}

func NewTxLoopUpsertOnlineStorage(db *postgres.Database) *TxLoopUpsertOnlineStorage {
	return &TxLoopUpsertOnlineStorage{db: db}
}

func (s *TxLoopUpsertOnlineStorage) BatchStore(ctx context.Context, pairs []UserOnlinePair) error {
	return s.db.WithTransaction(ctx, func(queries *dbs.Queries) error {
		for _, pair := range pairs {
			err := queries.UserOnlineUpsert(ctx, dbs.UserOnlineUpsertParams{
				UserID: pair.UserID,
				Online: pgtype.Timestamp{
					Time:  time.Unix(pair.Timestamp, 0).UTC(),
					Valid: true,
				},
			})
			if err != nil {
				return err
			}
		}

		return nil
	})
}

package main

import (
	"context"
	"testing"

	"github.com/doutivity/research-online-postgres-go/internal/storage/postgres"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/stretchr/testify/require"
)

func TestTxLoopUpsertOnlineStorage(t *testing.T) {
	ctx := context.Background()

	pool, err := pgxpool.New(ctx, dataSourceName)
	require.NoError(t, err)
	defer pool.Close()

	storage := NewTxLoopUpsertOnlineStorage(postgres.NewDatabase(pool))

	testOnlineStorage(t, storage)
}

package main

import (
	"context"

	"github.com/doutivity/research-online-postgres-go/internal/storage/postgres"
	"github.com/doutivity/research-online-postgres-go/internal/storage/postgres/dbs"
)

type UnnestUpsertOnlineStorage struct {
	db *postgres.Database
}

func NewUnnestUpsertOnlineStorage(db *postgres.Database) *UnnestUpsertOnlineStorage {
	return &UnnestUpsertOnlineStorage{db: db}
}

func (s *UnnestUpsertOnlineStorage) BatchStore(ctx context.Context, pairs []UserOnlinePair) error {
	userIDs, timestamps := userOnlinePairsToPgxSlices(pairs)

	return s.db.Queries().UserOnlineUnnestUpsert(ctx, dbs.UserOnlineUnnestUpsertParams{
		UserIds: userIDs,
		Onlines: timestamps,
	})
}

package main

import (
	"time"

	"github.com/jackc/pgx/v5/pgtype"
)

func userOnlinePairsToPgxSlices(pairs []UserOnlinePair) ([]int64, []pgtype.Timestamp) {
	var (
		userIDs    = make([]int64, len(pairs))
		timestamps = make([]pgtype.Timestamp, len(pairs))
	)

	for i, pair := range pairs {
		userIDs[i] = pair.UserID
		timestamps[i] = pgtype.Timestamp{
			Time:  time.Unix(pair.Timestamp, 0).UTC(),
			Valid: true,
		}
	}

	return userIDs, timestamps
}

package main

import (
	"context"
	"testing"

	"github.com/doutivity/research-online-postgres-go/internal/storage/postgres"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/stretchr/testify/require"
)

func TestUnnestUpsertOnlineStorage(t *testing.T) {
	ctx := context.Background()

	pool, err := pgxpool.New(ctx, dataSourceName)
	require.NoError(t, err)
	defer pool.Close()

	storage := NewUnnestUpsertOnlineStorage(postgres.NewDatabase(pool))

	testOnlineStorage(t, storage)
}

package main

import (
	"context"
	"sync/atomic"
	"testing"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/stretchr/testify/require"
)

func benchmarkOnlineStorage(
	b *testing.B,
	storage OnlineStorage,
) {
	b.Helper()

	const (
		batch  = 1000
		online = 1679800725
	)

	ctx := context.Background()

	pool, err := pgxpool.New(ctx, dataSourceName)
	require.NoError(b, err)
	defer pool.Close()

	truncateOnline(b, ctx, pool)
	generateOnline(b, ctx, pool, int64(b.N*batch), online)

	var (
		startTimestamp = time.Now().Unix()
		counter        = int64(0)
	)

	pairs := make([]UserOnlinePair, batch)

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		index := atomic.AddInt64(&counter, batch)

		for j := 0; j < batch; j++ {
			pairs[j] = UserOnlinePair{
				UserID:    int64(i*batch + j + 1), // 0 .. 99, 100 .. 199
				Timestamp: startTimestamp + index,
			}
		}

		err := storage.BatchStore(ctx, pairs)

		require.NoError(b, err)
	}
	b.StopTimer()

	expectedOnlineChangedCount(b, ctx, pool, int64(b.N*batch), online)
}

package main

import (
	"context"
	"testing"

	"github.com/doutivity/research-online-postgres-go/internal/storage/postgres"

	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/stretchr/testify/require"
)

func BenchmarkUnnestUpdateOnlineStorage(b *testing.B) {
	if testing.Short() {
		b.Skip()
	}

	ctx := context.Background()

	pool, err := pgxpool.New(ctx, dataSourceName)
	require.NoError(b, err)
	defer pool.Close()

	storage := NewUnnestUpdateOnlineStorage(postgres.NewDatabase(pool))

	benchmarkOnlineStorage(b, storage)
}