/
query.go
180 lines (152 loc) · 4.18 KB
/
query.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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
package server
import (
"context"
"fmt"
"sync"
"time"
"github.com/jackc/pgx/v5/pgproto3"
"github.com/mattermost/logr/v2"
)
const (
StatusUnset byte = 0
StatusIdle byte = 'I'
StatusInTx byte = 'T'
StatusError byte = 'E'
)
type ClientConn struct {
handle *pgproto3.Backend
txStatus byte
logger *logr.Logger
pool *Pool
// mut makes the getting of pool, and setting of the serverConn
// atomic. This allows us to find the serverConn for a clientConn
// to cancel a request.
mut sync.Mutex
// This is set to non-nil if there's an active transaction going on.
serverConn *ServerConn
schema string
}
func NewClientConn(handle *pgproto3.Backend, logger *logr.Logger, pool *Pool, schema string) *ClientConn {
return &ClientConn{
handle: handle,
logger: logger,
pool: pool,
schema: schema,
}
}
func (cc *ClientConn) handleQuery(feMsg pgproto3.FrontendMessage) error {
// Leasing a connection
if err := cc.acquireConn(); err != nil {
return err
}
serverEnd := pgproto3.NewFrontend(cc.serverConn.Conn(), cc.serverConn.Conn())
// serverEnd.Trace(cc.logger.Writer(), pgproto3.TracerOptions{})
serverEnd.Send(feMsg)
if err := serverEnd.Flush(); err != nil {
return fmt.Errorf("error while flushing queryMsg: %w", err)
}
if err := cc.readBackendResponse(serverEnd); err != nil {
return err
}
return nil
}
func (cc *ClientConn) handleExtendedQuery(feMsg pgproto3.FrontendMessage) error {
// Leasing a connection
if err := cc.acquireConn(); err != nil {
return err
}
serverEnd := pgproto3.NewFrontend(cc.serverConn.Conn(), cc.serverConn.Conn())
// serverEnd.Trace(cc.logger.Writer(), pgproto3.TracerOptions{})
serverEnd.Send(feMsg)
for {
feMsg, err := cc.handle.Receive()
if err != nil {
return fmt.Errorf("error while receiving msg in extendedQuery: %w", err)
}
serverEnd.Send(feMsg)
// Keep reading until we see a SYNC message
_, ok := feMsg.(*pgproto3.Sync)
if ok {
break
}
}
if err := serverEnd.Flush(); err != nil {
return fmt.Errorf("error while flushing extendedQuery: %w", err)
}
if err := cc.readBackendResponse(serverEnd); err != nil {
return err
}
return nil
}
func (cc *ClientConn) readBackendResponse(serverEnd *pgproto3.Frontend) error {
// Read the response
cnt := 0
for {
beMsg, err := serverEnd.Receive()
if err != nil {
return fmt.Errorf("error while receiving from server: %w", err)
}
cnt++
switch typedMsg := beMsg.(type) {
// Read all till ReadyForQuery
case *pgproto3.ReadyForQuery:
cc.handle.Send(typedMsg)
if err := cc.handle.Flush(); err != nil {
return fmt.Errorf("error while flushing to client: %w", err)
}
cc.txStatus = typedMsg.TxStatus
// Releasing the conn back to the pool
if cc.txStatus == StatusIdle {
cc.mut.Lock()
cc.pool.ReleaseConn(cc.serverConn)
cc.serverConn = nil
cc.mut.Unlock()
}
return nil
default:
cc.handle.Send(beMsg)
// Flush if we have queued too many messages
if cnt%10 == 0 {
if err := cc.handle.Flush(); err != nil {
return fmt.Errorf("error while flushing to client: %w", err)
}
}
continue
}
}
}
func (cc *ClientConn) acquireConn() error {
cc.mut.Lock()
defer cc.mut.Unlock()
if cc.serverConn != nil {
return nil
}
conn, err := cc.pool.AcquireConn()
if err != nil {
return fmt.Errorf("error while acquiring conn: %w", err)
}
// We have just got a connection from the pool. First, we check
// whether it's healthy or not.
if err2 := conn.CheckConn(); err2 != nil {
// *Server.handleConn will take care of closing the connection.
return fmt.Errorf("error while checking conn: %w", err2)
}
// This is a low-level method, so passing params is not really supported.
// We need to implement sanitization ourselves. XXX: item for future.
err = conn.Exec(fmt.Sprintf(`SET search_path='%s'`, cc.schema))
if err != nil {
return fmt.Errorf("error setting schema search path: %w", err)
}
cc.serverConn = conn
return nil
}
func (cc *ClientConn) CancelServerConn() error {
cc.mut.Lock()
if cc.serverConn == nil {
return nil
}
cc.mut.Unlock()
ctx, cancel := context.WithTimeout(context.Background(), time.Second*cc.pool.connCreateTimeout)
defer cancel()
return cc.serverConn.CancelRequest(ctx)
}