Skip to content

Commit 52888c5

Browse files
MoinTomrafiss
authored andcommitted
Add crdbpgxv5 for supporting github.com/jackc/pgx/v5
1 parent c97ddc6 commit 52888c5

File tree

6 files changed

+223
-17
lines changed

6 files changed

+223
-17
lines changed

crdb/crdbpgx/README.md

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
`crdbpgx` is a wrapper around the logic for issuing SQL transactions which
22
performs retries (as required by CockroachDB) when using
33
[`github.com/jackc/pgx`](https://github.com/jackc/pgx) in standalone-library
4-
mode. pgx versions below v4 are not supported.
4+
mode. pgx versions below v4 are not supported.
5+
6+
Note: use `crdbpgxv5` for `pgx/v5` support
57

68
If you're using pgx just as a driver for the standard `database/sql` package,
79
use the parent `crdb` package instead.

crdb/crdbpgxv5/README.md

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
`crdbpgxv5` is a wrapper around the logic for issuing SQL transactions which
2+
performs retries (as required by CockroachDB) when using
3+
[`github.com/jackc/pgx`](https://github.com/jackc/pgx) in standalone-library
4+
mode. `crdbpgxv5` only supports `pgx/v5`.
5+
6+
Note: use `crdbpgx` for `pgx/v4` support
7+
8+
If you're using pgx just as a driver for the standard `database/sql` package,
9+
use the parent `crdb` package instead.

crdb/crdbpgxv5/pgx.go

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
// Copyright 2022 The Cockroach Authors.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
12+
// implied. See the License for the specific language governing
13+
// permissions and limitations under the License.
14+
15+
package crdbpgx
16+
17+
import (
18+
"context"
19+
"github.com/jackc/pgx/v5"
20+
21+
"github.com/cockroachdb/cockroach-go/v2/crdb"
22+
)
23+
24+
// ExecuteTx runs fn inside a transaction and retries it as needed. On
25+
// non-retryable failures, the transaction is aborted and rolled back; on
26+
// success, the transaction is committed.
27+
//
28+
// See crdb.ExecuteTx() for more information.
29+
//
30+
// conn can be a pgx.Conn or a pgxpool.Pool.
31+
func ExecuteTx(
32+
ctx context.Context, conn Conn, txOptions pgx.TxOptions, fn func(pgx.Tx) error,
33+
) error {
34+
tx, err := conn.BeginTx(ctx, txOptions)
35+
if err != nil {
36+
return err
37+
}
38+
return crdb.ExecuteInTx(ctx, pgxTxAdapter{tx}, func() error { return fn(tx) })
39+
}
40+
41+
// Conn abstracts pgx transactions creators: pgx.Conn and pgxpool.Pool.
42+
type Conn interface {
43+
Begin(context.Context) (pgx.Tx, error)
44+
BeginTx(context.Context, pgx.TxOptions) (pgx.Tx, error)
45+
}
46+
47+
type pgxTxAdapter struct {
48+
tx pgx.Tx
49+
}
50+
51+
var _ crdb.Tx = pgxTxAdapter{}
52+
53+
func (tx pgxTxAdapter) Commit(ctx context.Context) error {
54+
return tx.tx.Commit(ctx)
55+
}
56+
57+
func (tx pgxTxAdapter) Rollback(ctx context.Context) error {
58+
return tx.tx.Rollback(ctx)
59+
}
60+
61+
// Exec is part of the crdb.Tx interface.
62+
func (tx pgxTxAdapter) Exec(ctx context.Context, q string, args ...interface{}) error {
63+
_, err := tx.tx.Exec(ctx, q, args...)
64+
return err
65+
}

crdb/crdbpgxv5/pgx_test.go

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
// Copyright 2022 The Cockroach Authors.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
12+
// implied. See the License for the specific language governing
13+
// permissions and limitations under the License.
14+
15+
package crdbpgx
16+
17+
import (
18+
"context"
19+
"fmt"
20+
"github.com/jackc/pgx/v5"
21+
"github.com/jackc/pgx/v5/pgxpool"
22+
"testing"
23+
24+
"github.com/cockroachdb/cockroach-go/v2/crdb"
25+
"github.com/cockroachdb/cockroach-go/v2/testserver"
26+
)
27+
28+
// TestExecuteTx verifies transaction retry using the classic
29+
// example of write skew in bank account balance transfers.
30+
func TestExecuteTx(t *testing.T) {
31+
ts, err := testserver.NewTestServer()
32+
if err != nil {
33+
t.Fatal(err)
34+
}
35+
ctx := context.Background()
36+
37+
pool, err := pgxpool.New(ctx, ts.PGURL().String())
38+
if err != nil {
39+
t.Fatal(err)
40+
}
41+
42+
if err := crdb.ExecuteTxGenericTest(ctx, pgxWriteSkewTest{pool: pool}); err != nil {
43+
t.Fatal(err)
44+
}
45+
}
46+
47+
type pgxWriteSkewTest struct {
48+
pool *pgxpool.Pool
49+
}
50+
51+
func (t pgxWriteSkewTest) Init(ctx context.Context) error {
52+
initStmt := `
53+
CREATE DATABASE d;
54+
CREATE TABLE d.t (acct INT PRIMARY KEY, balance INT);
55+
INSERT INTO d.t (acct, balance) VALUES (1, 100), (2, 100);
56+
`
57+
_, err := t.pool.Exec(ctx, initStmt)
58+
return err
59+
}
60+
61+
var _ crdb.WriteSkewTest = pgxWriteSkewTest{}
62+
63+
// ExecuteTx is part of the crdb.WriteSkewTest interface.
64+
func (t pgxWriteSkewTest) ExecuteTx(ctx context.Context, fn func(tx interface{}) error) error {
65+
return ExecuteTx(ctx, t.pool, pgx.TxOptions{}, func(tx pgx.Tx) error {
66+
return fn(tx)
67+
})
68+
}
69+
70+
// GetBalances is part of the crdb.WriteSkewTest interface.
71+
func (t pgxWriteSkewTest) GetBalances(ctx context.Context, txi interface{}) (int, int, error) {
72+
tx := txi.(pgx.Tx)
73+
var rows pgx.Rows
74+
rows, err := tx.Query(ctx, `SELECT balance FROM d.t WHERE acct IN (1, 2);`)
75+
if err != nil {
76+
return 0, 0, err
77+
}
78+
defer rows.Close()
79+
var bal1, bal2 int
80+
balances := []*int{&bal1, &bal2}
81+
i := 0
82+
for ; rows.Next(); i++ {
83+
if err = rows.Scan(balances[i]); err != nil {
84+
return 0, 0, err
85+
}
86+
}
87+
if i != 2 {
88+
return 0, 0, fmt.Errorf("expected two balances; got %d", i)
89+
}
90+
return bal1, bal2, nil
91+
}
92+
93+
// UpdateBalance is part of the crdb.WriteSkewInterface.
94+
func (t pgxWriteSkewTest) UpdateBalance(
95+
ctx context.Context, txi interface{}, acct, delta int,
96+
) error {
97+
tx := txi.(pgx.Tx)
98+
_, err := tx.Exec(ctx, `UPDATE d.t SET balance=balance+$1 WHERE acct=$2;`, delta, acct)
99+
return err
100+
}

go.mod

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,13 @@ go 1.13
55
require (
66
github.com/gofrs/flock v0.8.1
77
github.com/jackc/pgx/v4 v4.16.1
8+
github.com/jackc/pgx/v5 v5.2.0
89
github.com/jmoiron/sqlx v1.3.5
910
github.com/lib/pq v1.10.6
1011
github.com/pkg/errors v0.9.1 // indirect
1112
github.com/shopspring/decimal v1.3.1 // indirect
12-
github.com/stretchr/testify v1.7.0
13-
golang.org/x/crypto v0.0.0-20220517005047-85d78b3ac167 // indirect
14-
golang.org/x/sys v0.0.0-20220513210249-45d2b4557a2a // indirect
13+
github.com/stretchr/testify v1.8.1
1514
gorm.io/driver/postgres v1.3.5
1615
gorm.io/gorm v1.23.5
1716
)
17+

0 commit comments

Comments
 (0)