-
Notifications
You must be signed in to change notification settings - Fork 3
/
snapshotter.go
96 lines (80 loc) · 2.58 KB
/
snapshotter.go
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
package pglogicalstream
import (
"context"
"database/sql"
"fmt"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
_ "github.com/lib/pq"
"log"
"strings"
)
type Snapshotter struct {
pgConnection *pgx.Conn
snapshotName string
}
func NewSnapshotter(dbConf pgconn.Config, snapshotName string) (*Snapshotter, error) {
var sslMode = "none"
if dbConf.TLSConfig != nil {
sslMode = "require"
}
connStr := fmt.Sprintf("postgres://%s:%s@%s:%d/%s?sslmode=%s", dbConf.User,
dbConf.Password, dbConf.Host, dbConf.Port, dbConf.Database, sslMode,
)
pgConn, err := pgx.Connect(context.Background(), connStr)
return &Snapshotter{
pgConnection: pgConn,
snapshotName: snapshotName,
}, err
}
func (s *Snapshotter) Prepare() error {
if res, err := s.pgConnection.Exec(context.TODO(), "BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ;"); err != nil {
return err
} else {
fmt.Println(res.String())
}
if res, err := s.pgConnection.Exec(context.TODO(), fmt.Sprintf("SET TRANSACTION SNAPSHOT '%s';", s.snapshotName)); err != nil {
return err
} else {
fmt.Println(res.RowsAffected())
}
return nil
}
func (s *Snapshotter) FindAvgRowSize(table string) sql.NullInt64 {
var avgRowSize sql.NullInt64
if rows, err := s.pgConnection.Query(context.TODO(), fmt.Sprintf(`SELECT SUM(pg_column_size('%s.*')) / COUNT(*) FROM %s;`, table, table)); err != nil {
log.Fatal("Can get avg row size", err)
} else {
if rows.Next() {
if err = rows.Scan(&avgRowSize); err != nil {
log.Fatal("Can get avg row size", err)
}
} else {
log.Fatal("Can get avg row size; 0 rows returned")
}
}
return avgRowSize
}
func (s *Snapshotter) CalculateBatchSize(safetyFactor float64, availableMemory uint64, estimatedRowSize uint64) int {
// Adjust this factor based on your system's memory constraints.
// This example uses a safety factor of 0.8 to leave some memory headroom.
batchSize := int(float64(availableMemory) * safetyFactor / float64(estimatedRowSize))
if batchSize < 1 {
batchSize = 1
}
return batchSize
}
func (s *Snapshotter) QuerySnapshotData(table string, columns []string, pk string, limit, offset int) (rows pgx.Rows, err error) {
joinedColumns := strings.Join(columns, ", ")
return s.pgConnection.Query(context.TODO(), fmt.Sprintf("SELECT %s FROM %s ORDER BY %s LIMIT %d OFFSET %d;", joinedColumns, table, pk, limit, offset))
}
func (s *Snapshotter) ReleaseSnapshot() error {
_, err := s.pgConnection.Exec(context.TODO(), "COMMIT;")
return err
}
func (s *Snapshotter) CloseConn() error {
if s.pgConnection != nil {
return s.pgConnection.Close(context.TODO())
}
return nil
}