@@ -2,74 +2,146 @@ package pgx_test
22
33import (
44 "context"
5+ "errors"
6+ "fmt"
57 "os"
68 "testing"
9+ "time"
710
811 "github.com/jackc/pgx/v5"
9- "github.com/stretchr/testify/assert"
1012 "github.com/stretchr/testify/require"
13+ "golang.org/x/sync/errgroup"
1114)
1215
13- func TestPgbouncerStatementCacheDescribe (t * testing.T ) {
16+ func TestPgbouncerQueryExecModes (t * testing.T ) {
1417 connString := os .Getenv ("PGX_TEST_PGBOUNCER_CONN_STRING" )
1518 if connString == "" {
1619 t .Skipf ("Skipping due to missing environment variable %v" , "PGX_TEST_PGBOUNCER_CONN_STRING" )
1720 }
1821
19- config := mustParseConfig (t , connString )
20- config .DefaultQueryExecMode = pgx .QueryExecModeCacheDescribe
21- config .DescriptionCacheCapacity = 1024
22-
23- testPgbouncer (t , config , 10 , 100 )
24- }
25-
26- func TestPgbouncerSimpleProtocol (t * testing.T ) {
27- connString := os .Getenv ("PGX_TEST_PGBOUNCER_CONN_STRING" )
28- if connString == "" {
29- t .Skipf ("Skipping due to missing environment variable %v" , "PGX_TEST_PGBOUNCER_CONN_STRING" )
22+ tests := []struct {
23+ name string
24+ mode pgx.QueryExecMode
25+ statementCacheCapacity int
26+ descriptionCacheCapacity int
27+ }{
28+ {
29+ name : "cache statement" ,
30+ mode : pgx .QueryExecModeCacheStatement ,
31+ statementCacheCapacity : 32 ,
32+ },
33+ {
34+ name : "cache describe" ,
35+ mode : pgx .QueryExecModeCacheDescribe ,
36+ descriptionCacheCapacity : 32 ,
37+ },
38+ {
39+ name : "exec" ,
40+ mode : pgx .QueryExecModeExec ,
41+ },
42+ {
43+ name : "simple protocol" ,
44+ mode : pgx .QueryExecModeSimpleProtocol ,
45+ },
3046 }
3147
32- config := mustParseConfig (t , connString )
33- config .DefaultQueryExecMode = pgx .QueryExecModeSimpleProtocol
48+ for _ , tt := range tests {
49+ t .Run (tt .name , func (t * testing.T ) {
50+ config := mustParseConfig (t , connString )
51+ config .DefaultQueryExecMode = tt .mode
52+ config .StatementCacheCapacity = tt .statementCacheCapacity
53+ config .DescriptionCacheCapacity = tt .descriptionCacheCapacity
3454
35- testPgbouncer (t , config , 10 , 100 )
55+ testPgbouncer (t , config , 10 , 100 )
56+ })
57+ }
3658}
3759
3860func testPgbouncer (t * testing.T , config * pgx.ConnConfig , workers , iterations int ) {
39- doneChan := make (chan struct {})
61+ ctx , cancel := context .WithTimeout (t .Context (), 2 * time .Minute )
62+ defer cancel ()
63+
64+ eg , ctx := errgroup .WithContext (ctx )
4065
4166 for range workers {
42- go func () {
43- defer func () { doneChan <- struct {}{} }()
44- conn , err := pgx .ConnectConfig (context .Background (), config )
45- require .Nil (t , err )
46- defer closeConn (t , conn )
47-
48- for range iterations {
49- var i32 int32
50- var i64 int64
51- var f32 float32
52- var s string
53- var s2 string
54- err = conn .QueryRow (context .Background (), "select 1::int4, 2::int8, 3::float4, 'hi'::text" ).Scan (& i32 , & i64 , & f32 , & s )
55- require .NoError (t , err )
56- assert .Equal (t , int32 (1 ), i32 )
57- assert .Equal (t , int64 (2 ), i64 )
58- assert .Equal (t , float32 (3 ), f32 )
59- assert .Equal (t , "hi" , s )
60-
61- err = conn .QueryRow (context .Background (), "select 1::int8, 2::float4, 'bye'::text, 4::int4, 'whatever'::text" ).Scan (& i64 , & f32 , & s , & i32 , & s2 )
62- require .NoError (t , err )
63- assert .Equal (t , int64 (1 ), i64 )
64- assert .Equal (t , float32 (2 ), f32 )
65- assert .Equal (t , "bye" , s )
66- assert .Equal (t , int32 (4 ), i32 )
67- assert .Equal (t , "whatever" , s2 )
67+ eg .Go (func () (err error ) {
68+ conn , err := pgx .ConnectConfig (ctx , config )
69+ if err != nil {
70+ return err
6871 }
69- }()
72+ defer func () {
73+ err = errors .Join (err , conn .Close (ctx ))
74+ }()
75+
76+ return exercisePgbouncerConn (ctx , conn , iterations )
77+ })
7078 }
7179
72- for range workers {
73- <- doneChan
80+ require .NoError (t , eg .Wait ())
81+ }
82+
83+ func exercisePgbouncerConn (ctx context.Context , conn * pgx.Conn , iterations int ) error {
84+ for i := range iterations {
85+ var i32 int32
86+ var i64 int64
87+ var f32 float32
88+ var s string
89+
90+ err := conn .QueryRow (ctx , "select $1::int4, $2::int8, $3::float4, $4::text" , int32 (i ), int64 (i + 1 ), float32 (i + 2 ), "hi" ).Scan (& i32 , & i64 , & f32 , & s )
91+ if err != nil {
92+ return err
93+ }
94+ if i32 != int32 (i ) || i64 != int64 (i + 1 ) || f32 != float32 (i + 2 ) || s != "hi" {
95+ return fmt .Errorf ("unexpected query result: %d, %d, %f, %q" , i32 , i64 , f32 , s )
96+ }
97+ }
98+
99+ commandTag , err := conn .Exec (ctx , "select $1::int4" , int32 (42 ))
100+ if err != nil {
101+ return err
102+ }
103+ if commandTag .String () != "SELECT 1" {
104+ return fmt .Errorf ("unexpected command tag: %s" , commandTag )
105+ }
106+
107+ tx , err := conn .Begin (ctx )
108+ if err != nil {
109+ return err
110+ }
111+ var txValue int32
112+ if err := tx .QueryRow (ctx , "select $1::int4" , int32 (43 )).Scan (& txValue ); err != nil {
113+ _ = tx .Rollback (ctx )
114+ return err
115+ }
116+ if txValue != 43 {
117+ _ = tx .Rollback (ctx )
118+ return fmt .Errorf ("unexpected transaction query result: %d" , txValue )
119+ }
120+ if err := tx .Commit (ctx ); err != nil {
121+ return err
74122 }
123+
124+ batch := & pgx.Batch {}
125+ batch .Queue ("select $1::int4" , int32 (44 ))
126+ batch .Queue ("select $1::text" , "batch" )
127+ batchResults := conn .SendBatch (ctx , batch )
128+
129+ var batchInt int32
130+ if err := batchResults .QueryRow ().Scan (& batchInt ); err != nil {
131+ _ = batchResults .Close ()
132+ return err
133+ }
134+ var batchText string
135+ if err := batchResults .QueryRow ().Scan (& batchText ); err != nil {
136+ _ = batchResults .Close ()
137+ return err
138+ }
139+ if err := batchResults .Close (); err != nil {
140+ return err
141+ }
142+ if batchInt != 44 || batchText != "batch" {
143+ return fmt .Errorf ("unexpected batch query results: %d, %q" , batchInt , batchText )
144+ }
145+
146+ return nil
75147}
0 commit comments