txmng/txmng.go at main · slavaavr/txmng · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
package txmng

import (
	"context"
	"fmt"
)

type TxManager interface {
	RunTx(opts TxOpts, fn func(ctx TxContext) (Result, error)) (Result, error)
	RunNoTx(opts NoTxOpts, fn func(ctx NoTxContext) (Result, error)) (Result, error)
}

type DBManager[T any] interface {
	GetDB(ctx Context) (db T, rawCtx context.Context)
}

type manager[T any] struct {
	dbProvider        DBProvider[T]
	defaultFallbackDB T
	dynamicFallbackDB bool
}

func New[T any](p DBProvider[T], opts ...Option) (txm TxManager, dbm DBManager[T]) {
	var cfg config
	for _, opt := range opts {
		opt(&cfg)
	}

	m := &manager[T]{
		dbProvider:        p,
		defaultFallbackDB: p.GetDB(NoTxOpts{}),
		dynamicFallbackDB: cfg.dynamicFallbackDB,
	}

	txm, dbm = m, m
	if cfg.retrier != nil {
		txm = newTxManagerWithRetrier(txm, cfg.retrier)
	}

	return txm, dbm
}

func (s *manager[T]) RunTx(opts TxOpts, fn func(ctx TxContext) (Result, error)) (_ Result, err error) {
	if opts.Ctx == nil {
		opts.Ctx = context.Background()
	}

	tx, err := s.dbProvider.BeginTx(opts)
	if err != nil {
		return nil, fmt.Errorf("begin tx: %w", err)
	}

	fallbackDB := s.defaultFallbackDB
	if s.dynamicFallbackDB {
		fallbackDB = s.dbProvider.GetDB(newNoTxOpts(opts.Ctx, opts.Ext))
	}

	newCtx := newTxContext(opts.Ctx, tx.GetDB(), fallbackDB)

	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("panic: %v", r)
		}

		if err != nil {
			if err2 := tx.Rollback(opts.Ctx); err2 != nil {
				err = fmt.Errorf("rollback tx (%s): %w", err2, err)
			}
		}

		newCtx.close()
	}()

	res, err := fn(newCtx)
	if err != nil {
		return nil, err
	}

	if err = tx.Commit(opts.Ctx); err != nil {
		return nil, fmt.Errorf("commit tx: %w", err)
	}

	return res, nil
}

func (s *manager[T]) RunNoTx(opts NoTxOpts, fn func(ctx NoTxContext) (Result, error)) (_ Result, err error) {
	if opts.Ctx == nil {
		opts.Ctx = context.Background()
	}

	newCtx := newNoTxContext(opts.Ctx, s.dbProvider.GetDB(opts))

	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("panic: %v", r)
		}

		newCtx.close()
	}()

	return fn(newCtx)
}

func (s *manager[T]) GetDB(ctx Context) (db T, rawCtx context.Context) {
	return ctx.getDB().(T), ctx.rawCtx()
}