db-migrator.go

1
package connection
2

3
import (
4
	"context"
5
	"database/sql"
6
	"fmt"
7
	"strings"
8

9
	"github.com/pkg/errors"
10
)
11

12
type ContextKey string
13

14
const contextKeyTX ContextKey = "tx"
15

16
type Connection struct {
17
	driver Driver
18
	dsn    string
19
	db     *sql.DB
20
	ping   bool
21
}
22

23
func New(dsn string) (*Connection, error) {
24
	switch {
25
	case strings.HasPrefix(dsn, "clickhouse://"):
26
		return clickhouse(dsn)
27
	case strings.HasPrefix(dsn, "postgres://"):
28
		return postgres(dsn)
29
	case strings.HasPrefix(dsn, "mysql://"):
30
		return mysql(dsn)
31
	default:
32
		return nil, fmt.Errorf("driver \"%s\" doesn't support", dsn)
33
	}
34
}
35

36
// DSN returns the connection string.
37
func (c *Connection) DSN() string {
38
	return c.dsn
39
}
40

41
// Driver returns the driver name used to connect to the database.
42
func (c *Connection) Driver() Driver {
43
	return c.driver
44
}
45

46
// Ping checks connection
47
func (c *Connection) Ping() error {
48
	if c.ping {
49
		return nil
50
	}
51
	if err := c.db.Ping(); err != nil {
52
		return errors.Wrapf(err, "ping %v connection: %v", c.Driver(), c.dsn)
53
	}
54
	c.ping = true
55
	return nil
56
}
57

58
// QueryContext executes a query that returns rows, typically a SELECT.
59
// The args are for any placeholder parameters in the query.
60
func (c *Connection) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
61
	if err := c.Ping(); err != nil {
62
		return nil, err
63
	}
64
	return c.db.QueryContext(ctx, query, args...)
65
}
66

67
// ExecContext executes a query without returning any rows.
68
// The args are for any placeholder parameters in the query.
69
func (c *Connection) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
70
	if err := c.Ping(); err != nil {
71
		return nil, err
72
	}
73
	v := ctx.Value(contextKeyTX)
74
	if v != nil {
75
		if tx, ok := v.(*sql.Tx); ok {
76
			stmt, err := tx.PrepareContext(ctx, query)
77
			if err != nil {
78
				return nil, err
79
			}
80
			return stmt.ExecContext(ctx, args...)
81
		}
82
	}
83
	return c.db.ExecContext(ctx, query, args...)
84
}
85

86
// Transaction executes body in func txFn into transaction.
87
func (c *Connection) Transaction(ctx context.Context, txFn func(ctx context.Context) error) error {
88
	if err := c.Ping(); err != nil {
89
		return err
90
	}
91
	if v := ctx.Value(contextKeyTX); v != nil {
92
		return errors.New("active transaction does not close")
93
	}
94

95
	tx, err := c.db.BeginTx(ctx, nil)
96
	if err != nil {
97
		return err
98
	}
99

100
	ctxWithTX := context.WithValue(ctx, contextKeyTX, tx)
101

102
	if err := txFn(ctxWithTX); err != nil {
103
		if err2 := tx.Rollback(); err2 != nil {
104
			return errors.Wrap(err, err2.Error())
105
		}
106
		return err
107
	}
108

109
	return tx.Commit()
110
}
111

112
// clickhouse returns repository with clickhouse configuration.
113
func clickhouse(dsn string) (*Connection, error) {
114
	dsn, err := NormalizeClickhouseDSN(dsn)
115
	if err != nil {
116
		return nil, err
117
	}
118
	db, err := sql.Open("clickhouse", dsn)
119
	if err != nil {
120
		return nil, err
121
	}
122

123
	return &Connection{
124
		driver: DriverClickhouse,
125
		dsn:    dsn,
126
		db:     db,
127
	}, nil
128
}
129

130
// postgres returns repository with postgres configuration.
131
func postgres(dsn string) (*Connection, error) {
132
	db, err := sql.Open(DriverPostgres.String(), dsn)
133
	if err != nil {
134
		return nil, errors.Wrap(err, "open postgres connection")
135
	}
136

137
	return &Connection{
138
		driver: DriverPostgres,
139
		dsn:    dsn,
140
		db:     db,
141
	}, nil
142
}
143

144
// mysql returns repository with mysql configuration.
145
func mysql(dsn string) (*Connection, error) {
146
	db, err := sql.Open(DriverMySQL.String(), dsn[8:])
147
	if err != nil {
148
		return nil, errors.Wrap(err, "open mysql connection")
149
	}
150

151
	return &Connection{
152
		driver: DriverMySQL,
153
		dsn:    dsn,
154
		db:     db,
155
	}, nil
156
}
157